library/std/src/io/cursor.rs

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