1 #[cfg(test)]
2 mod tests;
3
4 use crate::std::io::prelude::*;
5
6 use crate::std::cmp;
7 use crate::std::io::{self, BorrowedCursor, ErrorKind, IoSlice, IoSliceMut, SeekFrom};
8 use core::alloc::Allocator;
9
10 /// A `Cursor` wraps an in-memory buffer and provides it with a
11 /// [`Seek`] implementation.
12 ///
13 /// `Cursor`s are used with in-memory buffers, anything implementing
14 /// <code>[AsRef]<\[u8]></code>, to allow them to implement [`Read`] and/or [`Write`],
15 /// allowing these buffers to be used anywhere you might use a reader or writer
16 /// that does actual I/O.
17 ///
18 /// The standard library implements some I/O traits on various types which
19 /// are commonly used as a buffer, like <code>Cursor<[Vec]\<u8>></code> and
20 /// <code>Cursor<[&\[u8\]][bytes]></code>.
21 ///
22 /// # Examples
23 ///
24 /// We may want to write bytes to a [`File`] in our production
25 /// code, but use an in-memory buffer in our tests. We can do this with
26 /// `Cursor`:
27 ///
28 /// [bytes]: crate::std::slice "slice"
29 /// [`File`]: crate::std::fs::File
30 ///
31 /// ```no_run
32 /// use std::io::prelude::*;
33 /// use std::io::{self, SeekFrom};
34 /// use std::fs::File;
35 ///
36 /// // a library function we've written
37 /// fn write_ten_bytes_at_end<W: Write + Seek>(mut writer: W) -> io::Result<()> {
38 /// writer.seek(SeekFrom::End(-10))?;
39 ///
40 /// for i in 0..10 {
41 /// writer.write(&[i])?;
42 /// }
43 ///
44 /// // all went well
45 /// Ok(())
46 /// }
47 ///
48 /// # fn foo() -> io::Result<()> {
49 /// // Here's some code that uses this library function.
50 /// //
51 /// // We might want to use a BufReader here for efficiency, but let's
52 /// // keep this example focused.
53 /// let mut file = File::create("foo.txt")?;
54 ///
55 /// write_ten_bytes_at_end(&mut file)?;
56 /// # Ok(())
57 /// # }
58 ///
59 /// // now let's write a test
60 /// #[test]
61 /// fn test_writes_bytes() {
62 /// // setting up a real File is much slower than an in-memory buffer,
63 /// // let's use a cursor instead
64 /// use std::io::Cursor;
65 /// let mut buff = Cursor::new(vec![0; 15]);
66 ///
67 /// write_ten_bytes_at_end(&mut buff).unwrap();
68 ///
69 /// assert_eq!(&buff.get_ref()[5..15], &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
70 /// }
71 /// ```
72 #[derive(Debug, Default, Eq, PartialEq)]
73 pub struct Cursor<T> {
74 inner: T,
75 pos: u64,
76 }
77
78 impl<T> Cursor<T> {
79 /// Creates a new cursor wrapping the provided underlying in-memory buffer.
80 ///
81 /// Cursor initial position is `0` even if underlying buffer (e.g., [`Vec`])
82 /// is not empty. So writing to cursor starts with overwriting [`Vec`]
83 /// content, not with appending to it.
84 ///
85 /// # Examples
86 ///
87 /// ```
88 /// use std::io::Cursor;
89 ///
90 /// let buff = Cursor::new(Vec::new());
91 /// # fn force_inference(_: &Cursor<Vec<u8>>) {}
92 /// # force_inference(&buff);
93 /// ```
new(inner: T) -> Cursor<T>94 pub const fn new(inner: T) -> Cursor<T> {
95 Cursor { pos: 0, inner }
96 }
97
98 /// Consumes this cursor, returning the underlying value.
99 ///
100 /// # Examples
101 ///
102 /// ```
103 /// use std::io::Cursor;
104 ///
105 /// let buff = Cursor::new(Vec::new());
106 /// # fn force_inference(_: &Cursor<Vec<u8>>) {}
107 /// # force_inference(&buff);
108 ///
109 /// let vec = buff.into_inner();
110 /// ```
into_inner(self) -> T111 pub fn into_inner(self) -> T {
112 self.inner
113 }
114
115 /// Gets a reference to the underlying value in this cursor.
116 ///
117 /// # Examples
118 ///
119 /// ```
120 /// use std::io::Cursor;
121 ///
122 /// let buff = Cursor::new(Vec::new());
123 /// # fn force_inference(_: &Cursor<Vec<u8>>) {}
124 /// # force_inference(&buff);
125 ///
126 /// let reference = buff.get_ref();
127 /// ```
get_ref(&self) -> &T128 pub const fn get_ref(&self) -> &T {
129 &self.inner
130 }
131
132 /// Gets a mutable reference to the underlying value in this cursor.
133 ///
134 /// Care should be taken to avoid modifying the internal I/O state of the
135 /// underlying value as it may corrupt this cursor's position.
136 ///
137 /// # Examples
138 ///
139 /// ```
140 /// use std::io::Cursor;
141 ///
142 /// let mut buff = Cursor::new(Vec::new());
143 /// # fn force_inference(_: &Cursor<Vec<u8>>) {}
144 /// # force_inference(&buff);
145 ///
146 /// let reference = buff.get_mut();
147 /// ```
get_mut(&mut self) -> &mut T148 pub fn get_mut(&mut self) -> &mut T {
149 &mut self.inner
150 }
151
152 /// Returns the current position of this cursor.
153 ///
154 /// # Examples
155 ///
156 /// ```
157 /// use std::io::Cursor;
158 /// use std::io::prelude::*;
159 /// use std::io::SeekFrom;
160 ///
161 /// let mut buff = Cursor::new(vec![1, 2, 3, 4, 5]);
162 ///
163 /// assert_eq!(buff.position(), 0);
164 ///
165 /// buff.seek(SeekFrom::Current(2)).unwrap();
166 /// assert_eq!(buff.position(), 2);
167 ///
168 /// buff.seek(SeekFrom::Current(-1)).unwrap();
169 /// assert_eq!(buff.position(), 1);
170 /// ```
position(&self) -> u64171 pub const fn position(&self) -> u64 {
172 self.pos
173 }
174
175 /// Sets the position of this cursor.
176 ///
177 /// # Examples
178 ///
179 /// ```
180 /// use std::io::Cursor;
181 ///
182 /// let mut buff = Cursor::new(vec![1, 2, 3, 4, 5]);
183 ///
184 /// assert_eq!(buff.position(), 0);
185 ///
186 /// buff.set_position(2);
187 /// assert_eq!(buff.position(), 2);
188 ///
189 /// buff.set_position(4);
190 /// assert_eq!(buff.position(), 4);
191 /// ```
set_position(&mut self, pos: u64)192 pub fn set_position(&mut self, pos: u64) {
193 self.pos = pos;
194 }
195 }
196
197 impl<T> Cursor<T>
198 where
199 T: AsRef<[u8]>,
200 {
201 /// Returns the remaining slice.
202 ///
203 /// # Examples
204 ///
205 /// ```
206 /// #![feature(cursor_remaining)]
207 /// use std::io::Cursor;
208 ///
209 /// let mut buff = Cursor::new(vec![1, 2, 3, 4, 5]);
210 ///
211 /// assert_eq!(buff.remaining_slice(), &[1, 2, 3, 4, 5]);
212 ///
213 /// buff.set_position(2);
214 /// assert_eq!(buff.remaining_slice(), &[3, 4, 5]);
215 ///
216 /// buff.set_position(4);
217 /// assert_eq!(buff.remaining_slice(), &[5]);
218 ///
219 /// buff.set_position(6);
220 /// assert_eq!(buff.remaining_slice(), &[]);
221 /// ```
remaining_slice(&self) -> &[u8]222 pub fn remaining_slice(&self) -> &[u8] {
223 let len = self.pos.min(self.inner.as_ref().len() as u64);
224 &self.inner.as_ref()[(len as usize)..]
225 }
226
227 /// Returns `true` if the remaining slice is empty.
228 ///
229 /// # Examples
230 ///
231 /// ```
232 /// #![feature(cursor_remaining)]
233 /// use std::io::Cursor;
234 ///
235 /// let mut buff = Cursor::new(vec![1, 2, 3, 4, 5]);
236 ///
237 /// buff.set_position(2);
238 /// assert!(!buff.is_empty());
239 ///
240 /// buff.set_position(5);
241 /// assert!(buff.is_empty());
242 ///
243 /// buff.set_position(10);
244 /// assert!(buff.is_empty());
245 /// ```
is_empty(&self) -> bool246 pub fn is_empty(&self) -> bool {
247 self.pos >= self.inner.as_ref().len() as u64
248 }
249 }
250
251 impl<T> Clone for Cursor<T>
252 where
253 T: Clone,
254 {
255 #[inline]
clone(&self) -> Self256 fn clone(&self) -> Self {
257 Cursor {
258 inner: self.inner.clone(),
259 pos: self.pos,
260 }
261 }
262
263 #[inline]
clone_from(&mut self, other: &Self)264 fn clone_from(&mut self, other: &Self) {
265 self.inner.clone_from(&other.inner);
266 self.pos = other.pos;
267 }
268 }
269
270 impl<T> io::Seek for Cursor<T>
271 where
272 T: AsRef<[u8]>,
273 {
seek(&mut self, style: SeekFrom) -> io::Result<u64>274 fn seek(&mut self, style: SeekFrom) -> io::Result<u64> {
275 let (base_pos, offset) = match style {
276 SeekFrom::Start(n) => {
277 self.pos = n;
278 return Ok(n);
279 }
280 SeekFrom::End(n) => (self.inner.as_ref().len() as u64, n),
281 SeekFrom::Current(n) => (self.pos, n),
282 };
283 match base_pos.checked_add_signed(offset) {
284 Some(n) => {
285 self.pos = n;
286 Ok(self.pos)
287 }
288 None => Err(io::const_io_error!(
289 ErrorKind::InvalidInput,
290 "invalid seek to a negative or overflowing position",
291 )),
292 }
293 }
294
stream_len(&mut self) -> io::Result<u64>295 fn stream_len(&mut self) -> io::Result<u64> {
296 Ok(self.inner.as_ref().len() as u64)
297 }
298
stream_position(&mut self) -> io::Result<u64>299 fn stream_position(&mut self) -> io::Result<u64> {
300 Ok(self.pos)
301 }
302 }
303
304 impl<T> Read for Cursor<T>
305 where
306 T: AsRef<[u8]>,
307 {
read(&mut self, buf: &mut [u8]) -> io::Result<usize>308 fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
309 let n = Read::read(&mut self.remaining_slice(), buf)?;
310 self.pos += n as u64;
311 Ok(n)
312 }
313
read_buf(&mut self, mut cursor: BorrowedCursor<'_>) -> io::Result<()>314 fn read_buf(&mut self, mut cursor: BorrowedCursor<'_>) -> io::Result<()> {
315 let prev_written = cursor.written();
316
317 Read::read_buf(&mut self.fill_buf()?, cursor.reborrow())?;
318
319 self.pos += (cursor.written() - prev_written) as u64;
320
321 Ok(())
322 }
323
read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize>324 fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
325 let mut nread = 0;
326 for buf in bufs {
327 let n = self.read(buf)?;
328 nread += n;
329 if n < buf.len() {
330 break;
331 }
332 }
333 Ok(nread)
334 }
335
is_read_vectored(&self) -> bool336 fn is_read_vectored(&self) -> bool {
337 true
338 }
339
read_exact(&mut self, buf: &mut [u8]) -> io::Result<()>340 fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
341 let n = buf.len();
342 Read::read_exact(&mut self.remaining_slice(), buf)?;
343 self.pos += n as u64;
344 Ok(())
345 }
346 }
347
348 impl<T> BufRead for Cursor<T>
349 where
350 T: AsRef<[u8]>,
351 {
fill_buf(&mut self) -> io::Result<&[u8]>352 fn fill_buf(&mut self) -> io::Result<&[u8]> {
353 Ok(self.remaining_slice())
354 }
consume(&mut self, amt: usize)355 fn consume(&mut self, amt: usize) {
356 self.pos += amt as u64;
357 }
358 }
359
360 // Non-resizing write implementation
361 #[inline]
slice_write(pos_mut: &mut u64, slice: &mut [u8], buf: &[u8]) -> io::Result<usize>362 fn slice_write(pos_mut: &mut u64, slice: &mut [u8], buf: &[u8]) -> io::Result<usize> {
363 let pos = cmp::min(*pos_mut, slice.len() as u64);
364 let amt = (&mut slice[(pos as usize)..]).write(buf)?;
365 *pos_mut += amt as u64;
366 Ok(amt)
367 }
368
369 #[inline]
slice_write_vectored( pos_mut: &mut u64, slice: &mut [u8], bufs: &[IoSlice<'_>], ) -> io::Result<usize>370 fn slice_write_vectored(
371 pos_mut: &mut u64,
372 slice: &mut [u8],
373 bufs: &[IoSlice<'_>],
374 ) -> io::Result<usize> {
375 let mut nwritten = 0;
376 for buf in bufs {
377 let n = slice_write(pos_mut, slice, buf)?;
378 nwritten += n;
379 if n < buf.len() {
380 break;
381 }
382 }
383 Ok(nwritten)
384 }
385
386 /// Reserves the required space, and pads the vec with 0s if necessary.
reserve_and_pad<A: Allocator>( pos_mut: &mut u64, vec: &mut Vec<u8, A>, buf_len: usize, ) -> io::Result<usize>387 fn reserve_and_pad<A: Allocator>(
388 pos_mut: &mut u64,
389 vec: &mut Vec<u8, A>,
390 buf_len: usize,
391 ) -> io::Result<usize> {
392 let pos: usize = (*pos_mut).try_into().map_err(|_| {
393 io::const_io_error!(
394 ErrorKind::InvalidInput,
395 "cursor position exceeds maximum possible vector length",
396 )
397 })?;
398
399 // For safety reasons, we don't want these numbers to overflow
400 // otherwise our allocation won't be enough
401 let desired_cap = pos.saturating_add(buf_len);
402 if desired_cap > vec.capacity() {
403 // We want our vec's total capacity
404 // to have room for (pos+buf_len) bytes. Reserve allocates
405 // based on additional elements from the length, so we need to
406 // reserve the difference
407 vec.reserve(desired_cap - vec.len());
408 }
409 // Pad if pos is above the current len.
410 if pos > vec.len() {
411 let diff = pos - vec.len();
412 // Unfortunately, `resize()` would suffice but the optimiser does not
413 // realise the `reserve` it does can be eliminated. So we do it manually
414 // to eliminate that extra branch
415 let spare = vec.spare_capacity_mut();
416 debug_assert!(spare.len() >= diff);
417 // Safety: we have allocated enough capacity for this.
418 // And we are only writing, not reading
419 unsafe {
420 spare
421 .get_unchecked_mut(..diff)
422 .fill(core::mem::MaybeUninit::new(0));
423 vec.set_len(pos);
424 }
425 }
426
427 Ok(pos)
428 }
429
430 /// Writes the slice to the vec without allocating
431 /// # Safety: vec must have buf.len() spare capacity
vec_write_unchecked<A>(pos: usize, vec: &mut Vec<u8, A>, buf: &[u8]) -> usize where A: Allocator,432 unsafe fn vec_write_unchecked<A>(pos: usize, vec: &mut Vec<u8, A>, buf: &[u8]) -> usize
433 where
434 A: Allocator,
435 {
436 debug_assert!(vec.capacity() >= pos + buf.len());
437 vec.as_mut_ptr().add(pos).copy_from(buf.as_ptr(), buf.len());
438 pos + buf.len()
439 }
440
441 /// Resizing write implementation for [`Cursor`]
442 ///
443 /// Cursor is allowed to have a pre-allocated and initialised
444 /// vector body, but with a position of 0. This means the [`Write`]
445 /// will overwrite the contents of the vec.
446 ///
447 /// This also allows for the vec body to be empty, but with a position of N.
448 /// This means that [`Write`] will pad the vec with 0 initially,
449 /// before writing anything from that point
vec_write<A>(pos_mut: &mut u64, vec: &mut Vec<u8, A>, buf: &[u8]) -> io::Result<usize> where A: Allocator,450 fn vec_write<A>(pos_mut: &mut u64, vec: &mut Vec<u8, A>, buf: &[u8]) -> io::Result<usize>
451 where
452 A: Allocator,
453 {
454 let buf_len = buf.len();
455 let mut pos = reserve_and_pad(pos_mut, vec, buf_len)?;
456
457 // Write the buf then progress the vec forward if necessary
458 // Safety: we have ensured that the capacity is available
459 // and that all bytes get written up to pos
460 unsafe {
461 pos = vec_write_unchecked(pos, vec, buf);
462 if pos > vec.len() {
463 vec.set_len(pos);
464 }
465 };
466
467 // Bump us forward
468 *pos_mut += buf_len as u64;
469 Ok(buf_len)
470 }
471
472 /// Resizing write_vectored implementation for [`Cursor`]
473 ///
474 /// Cursor is allowed to have a pre-allocated and initialised
475 /// vector body, but with a position of 0. This means the [`Write`]
476 /// will overwrite the contents of the vec.
477 ///
478 /// This also allows for the vec body to be empty, but with a position of N.
479 /// This means that [`Write`] will pad the vec with 0 initially,
480 /// before writing anything from that point
vec_write_vectored<A>( pos_mut: &mut u64, vec: &mut Vec<u8, A>, bufs: &[IoSlice<'_>], ) -> io::Result<usize> where A: Allocator,481 fn vec_write_vectored<A>(
482 pos_mut: &mut u64,
483 vec: &mut Vec<u8, A>,
484 bufs: &[IoSlice<'_>],
485 ) -> io::Result<usize>
486 where
487 A: Allocator,
488 {
489 // For safety reasons, we don't want this sum to overflow ever.
490 // If this saturates, the reserve should panic to avoid any unsound writing.
491 let buf_len = bufs.iter().fold(0usize, |a, b| a.saturating_add(b.len()));
492 let mut pos = reserve_and_pad(pos_mut, vec, buf_len)?;
493
494 // Write the buf then progress the vec forward if necessary
495 // Safety: we have ensured that the capacity is available
496 // and that all bytes get written up to the last pos
497 unsafe {
498 for buf in bufs {
499 pos = vec_write_unchecked(pos, vec, buf);
500 }
501 if pos > vec.len() {
502 vec.set_len(pos);
503 }
504 }
505
506 // Bump us forward
507 *pos_mut += buf_len as u64;
508 Ok(buf_len)
509 }
510
511 impl Write for Cursor<&mut [u8]> {
512 #[inline]
write(&mut self, buf: &[u8]) -> io::Result<usize>513 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
514 slice_write(&mut self.pos, self.inner, buf)
515 }
516
517 #[inline]
write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize>518 fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
519 slice_write_vectored(&mut self.pos, self.inner, bufs)
520 }
521
522 #[inline]
is_write_vectored(&self) -> bool523 fn is_write_vectored(&self) -> bool {
524 true
525 }
526
527 #[inline]
flush(&mut self) -> io::Result<()>528 fn flush(&mut self) -> io::Result<()> {
529 Ok(())
530 }
531 }
532
533 impl<A> Write for Cursor<&mut Vec<u8, A>>
534 where
535 A: Allocator,
536 {
write(&mut self, buf: &[u8]) -> io::Result<usize>537 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
538 vec_write(&mut self.pos, self.inner, buf)
539 }
540
write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize>541 fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
542 vec_write_vectored(&mut self.pos, self.inner, bufs)
543 }
544
545 #[inline]
is_write_vectored(&self) -> bool546 fn is_write_vectored(&self) -> bool {
547 true
548 }
549
550 #[inline]
flush(&mut self) -> io::Result<()>551 fn flush(&mut self) -> io::Result<()> {
552 Ok(())
553 }
554 }
555
556 impl<A> Write for Cursor<Vec<u8, A>>
557 where
558 A: Allocator,
559 {
write(&mut self, buf: &[u8]) -> io::Result<usize>560 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
561 vec_write(&mut self.pos, &mut self.inner, buf)
562 }
563
write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize>564 fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
565 vec_write_vectored(&mut self.pos, &mut self.inner, bufs)
566 }
567
568 #[inline]
is_write_vectored(&self) -> bool569 fn is_write_vectored(&self) -> bool {
570 true
571 }
572
573 #[inline]
flush(&mut self) -> io::Result<()>574 fn flush(&mut self) -> io::Result<()> {
575 Ok(())
576 }
577 }
578
579 impl<A> Write for Cursor<Box<[u8], A>>
580 where
581 A: Allocator,
582 {
583 #[inline]
write(&mut self, buf: &[u8]) -> io::Result<usize>584 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
585 slice_write(&mut self.pos, &mut self.inner, buf)
586 }
587
588 #[inline]
write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize>589 fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
590 slice_write_vectored(&mut self.pos, &mut self.inner, bufs)
591 }
592
593 #[inline]
is_write_vectored(&self) -> bool594 fn is_write_vectored(&self) -> bool {
595 true
596 }
597
598 #[inline]
flush(&mut self) -> io::Result<()>599 fn flush(&mut self) -> io::Result<()> {
600 Ok(())
601 }
602 }
603
604 impl<const N: usize> Write for Cursor<[u8; N]> {
605 #[inline]
write(&mut self, buf: &[u8]) -> io::Result<usize>606 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
607 slice_write(&mut self.pos, &mut self.inner, buf)
608 }
609
610 #[inline]
write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize>611 fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
612 slice_write_vectored(&mut self.pos, &mut self.inner, bufs)
613 }
614
615 #[inline]
is_write_vectored(&self) -> bool616 fn is_write_vectored(&self) -> bool {
617 true
618 }
619
620 #[inline]
flush(&mut self) -> io::Result<()>621 fn flush(&mut self) -> io::Result<()> {
622 Ok(())
623 }
624 }
625