src/libstd/io/buffered.rs

   1 // Copyright 2013 The Rust Project Developers. See the COPYRIGHT
   2 // file at the top-level directory of this distribution and at
   3 // http://rust-lang.org/COPYRIGHT.
   4 //
   5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
   6 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
   7 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
   8 // option. This file may not be copied, modified, or distributed
   9 // except according to those terms.
  10
  11 //! Buffering wrappers for I/O traits
  12
  13 use cmp;
  14 use container::Container;
  15 use io::{Reader, Writer, Stream, Buffer, DEFAULT_BUF_SIZE, IoResult};
  16 use iter::ExactSize;
  17 use ops::Drop;
  18 use option::{Some, None, Option};
  19 use result::{Ok, Err};
  20 use slice::{ImmutableVector, MutableVector};
  21 use slice;
  22 use vec::Vec;
  23
  24 /// Wraps a Reader and buffers input from it
  25 ///
  26 /// It can be excessively inefficient to work directly with a `Reader`. For
  27 /// example, every call to `read` on `TcpStream` results in a system call. A
  28 /// `BufferedReader` performs large, infrequent reads on the underlying
  29 /// `Reader` and maintains an in-memory buffer of the results.
  30 ///
  31 /// # Example
  32 ///
  33 /// ```rust
  34 /// use std::io::{BufferedReader, File};
  35 ///
  36 /// let file = File::open(&Path::new("message.txt"));
  37 /// let mut reader = BufferedReader::new(file);
  38 ///
  39 /// let mut buf = [0, ..100];
  40 /// match reader.read(buf) {
  41 ///     Ok(nread) => println!("Read {} bytes", nread),
  42 ///     Err(e) => println!("error reading: {}", e)
  43 /// }
  44 /// ```
  45 pub struct BufferedReader<R> {
  46     inner: R,
  47     buf: Vec<u8>,
  48     pos: uint,
  49     cap: uint,
  50 }
  51
  52 impl<R: Reader> BufferedReader<R> {
  53     /// Creates a new `BufferedReader` with the specified buffer capacity
  54     pub fn with_capacity(cap: uint, inner: R) -> BufferedReader<R> {
  55         // It's *much* faster to create an uninitialized buffer than it is to
  56         // fill everything in with 0. This buffer is entirely an implementation
  57         // detail and is never exposed, so we're safe to not initialize
  58         // everything up-front. This allows creation of BufferedReader instances
  59         // to be very cheap (large mallocs are not nearly as expensive as large
  60         // callocs).
  61         let mut buf = Vec::with_capacity(cap);
  62         unsafe { buf.set_len(cap); }
  63         BufferedReader {
  64             inner: inner,
  65             buf: buf,
  66             pos: 0,
  67             cap: 0,
  68         }
  69     }
  70
  71     /// Creates a new `BufferedReader` with a default buffer capacity
  72     pub fn new(inner: R) -> BufferedReader<R> {
  73         BufferedReader::with_capacity(DEFAULT_BUF_SIZE, inner)
  74     }
  75
  76     /// Gets a reference to the underlying reader.
  77     ///
  78     /// This type does not expose the ability to get a mutable reference to the
  79     /// underlying reader because that could possibly corrupt the buffer.
  80     pub fn get_ref<'a>(&'a self) -> &'a R { &self.inner }
  81
  82     /// Unwraps this `BufferedReader`, returning the underlying reader.
  83     ///
  84     /// Note that any leftover data in the internal buffer is lost.
  85     pub fn unwrap(self) -> R { self.inner }
  86 }
  87
  88 impl<R: Reader> Buffer for BufferedReader<R> {
  89     fn fill_buf<'a>(&'a mut self) -> IoResult<&'a [u8]> {
  90         if self.pos == self.cap {
  91             self.cap = try!(self.inner.read(self.buf.as_mut_slice()));
  92             self.pos = 0;
  93         }
  94         Ok(self.buf.slice(self.pos, self.cap))
  95     }
  96
  97     fn consume(&mut self, amt: uint) {
  98         self.pos += amt;
  99         assert!(self.pos <= self.cap);
 100     }
 101 }
 102
 103 impl<R: Reader> Reader for BufferedReader<R> {
 104     fn read(&mut self, buf: &mut [u8]) -> IoResult<uint> {
 105         let nread = {
 106             let available = try!(self.fill_buf());
 107             let nread = cmp::min(available.len(), buf.len());
 108             slice::bytes::copy_memory(buf, available.slice_to(nread));
 109             nread
 110         };
 111         self.pos += nread;
 112         Ok(nread)
 113     }
 114 }
 115
 116 /// Wraps a Writer and buffers output to it
 117 ///
 118 /// It can be excessively inefficient to work directly with a `Writer`. For
 119 /// example, every call to `write` on `TcpStream` results in a system call. A
 120 /// `BufferedWriter` keeps an in memory buffer of data and writes it to the
 121 /// underlying `Writer` in large, infrequent batches.
 122 ///
 123 /// This writer will be flushed when it is dropped.
 124 ///
 125 /// # Example
 126 ///
 127 /// ```rust
 128 /// # #![allow(unused_must_use)]
 129 /// use std::io::{BufferedWriter, File};
 130 ///
 131 /// let file = File::open(&Path::new("message.txt"));
 132 /// let mut writer = BufferedWriter::new(file);
 133 ///
 134 /// writer.write_str("hello, world");
 135 /// writer.flush();
 136 /// ```
 137 pub struct BufferedWriter<W> {
 138     inner: Option<W>,
 139     buf: Vec<u8>,
 140     pos: uint
 141 }
 142
 143 impl<W: Writer> BufferedWriter<W> {
 144     /// Creates a new `BufferedWriter` with the specified buffer capacity
 145     pub fn with_capacity(cap: uint, inner: W) -> BufferedWriter<W> {
 146         // See comments in BufferedReader for why this uses unsafe code.
 147         let mut buf = Vec::with_capacity(cap);
 148         unsafe { buf.set_len(cap); }
 149         BufferedWriter {
 150             inner: Some(inner),
 151             buf: buf,
 152             pos: 0
 153         }
 154     }
 155
 156     /// Creates a new `BufferedWriter` with a default buffer capacity
 157     pub fn new(inner: W) -> BufferedWriter<W> {
 158         BufferedWriter::with_capacity(DEFAULT_BUF_SIZE, inner)
 159     }
 160
 161     fn flush_buf(&mut self) -> IoResult<()> {
 162         if self.pos != 0 {
 163             let ret = self.inner.get_mut_ref().write(self.buf.slice_to(self.pos));
 164             self.pos = 0;
 165             ret
 166         } else {
 167             Ok(())
 168         }
 169     }
 170
 171     /// Gets a reference to the underlying writer.
 172     ///
 173     /// This type does not expose the ability to get a mutable reference to the
 174     /// underlying reader because that could possibly corrupt the buffer.
 175     pub fn get_ref<'a>(&'a self) -> &'a W { self.inner.get_ref() }
 176
 177     /// Unwraps this `BufferedWriter`, returning the underlying writer.
 178     ///
 179     /// The buffer is flushed before returning the writer.
 180     pub fn unwrap(mut self) -> W {
 181         // FIXME(#12628): is failing the right thing to do if flushing fails?
 182         self.flush_buf().unwrap();
 183         self.inner.take_unwrap()
 184     }
 185 }
 186
 187 impl<W: Writer> Writer for BufferedWriter<W> {
 188     fn write(&mut self, buf: &[u8]) -> IoResult<()> {
 189         if self.pos + buf.len() > self.buf.len() {
 190             try!(self.flush_buf());
 191         }
 192
 193         if buf.len() > self.buf.len() {
 194             self.inner.get_mut_ref().write(buf)
 195         } else {
 196             let dst = self.buf.mut_slice_from(self.pos);
 197             slice::bytes::copy_memory(dst, buf);
 198             self.pos += buf.len();
 199             Ok(())
 200         }
 201     }
 202
 203     fn flush(&mut self) -> IoResult<()> {
 204         self.flush_buf().and_then(|()| self.inner.get_mut_ref().flush())
 205     }
 206 }
 207
 208 #[unsafe_destructor]
 209 impl<W: Writer> Drop for BufferedWriter<W> {
 210     fn drop(&mut self) {
 211         if self.inner.is_some() {
 212             // FIXME(#12628): should this error be ignored?
 213             let _ = self.flush_buf();
 214         }
 215     }
 216 }
 217
 218 /// Wraps a Writer and buffers output to it, flushing whenever a newline (`0x0a`,
 219 /// `'\n'`) is detected.
 220 ///
 221 /// This writer will be flushed when it is dropped.
 222 pub struct LineBufferedWriter<W> {
 223     inner: BufferedWriter<W>,
 224 }
 225
 226 impl<W: Writer> LineBufferedWriter<W> {
 227     /// Creates a new `LineBufferedWriter`
 228     pub fn new(inner: W) -> LineBufferedWriter<W> {
 229         // Lines typically aren't that long, don't use a giant buffer
 230         LineBufferedWriter {
 231             inner: BufferedWriter::with_capacity(1024, inner)
 232         }
 233     }
 234
 235     /// Gets a reference to the underlying writer.
 236     ///
 237     /// This type does not expose the ability to get a mutable reference to the
 238     /// underlying reader because that could possibly corrupt the buffer.
 239     pub fn get_ref<'a>(&'a self) -> &'a W { self.inner.get_ref() }
 240
 241     /// Unwraps this `LineBufferedWriter`, returning the underlying writer.
 242     ///
 243     /// The internal buffer is flushed before returning the writer.
 244     pub fn unwrap(self) -> W { self.inner.unwrap() }
 245 }
 246
 247 impl<W: Writer> Writer for LineBufferedWriter<W> {
 248     fn write(&mut self, buf: &[u8]) -> IoResult<()> {
 249         match buf.iter().rposition(|&b| b == '\n' as u8) {
 250             Some(i) => {
 251                 try!(self.inner.write(buf.slice_to(i + 1)));
 252                 try!(self.inner.flush());
 253                 try!(self.inner.write(buf.slice_from(i + 1)));
 254                 Ok(())
 255             }
 256             None => self.inner.write(buf),
 257         }
 258     }
 259
 260     fn flush(&mut self) -> IoResult<()> { self.inner.flush() }
 261 }
 262
 263 struct InternalBufferedWriter<W>(BufferedWriter<W>);
 264
 265 impl<W> InternalBufferedWriter<W> {
 266     fn get_mut_ref<'a>(&'a mut self) -> &'a mut BufferedWriter<W> {
 267         let InternalBufferedWriter(ref mut w) = *self;
 268         return w;
 269     }
 270 }
 271
 272 impl<W: Reader> Reader for InternalBufferedWriter<W> {
 273     fn read(&mut self, buf: &mut [u8]) -> IoResult<uint> {
 274         self.get_mut_ref().inner.get_mut_ref().read(buf)
 275     }
 276 }
 277
 278 /// Wraps a Stream and buffers input and output to and from it.
 279 ///
 280 /// It can be excessively inefficient to work directly with a `Stream`. For
 281 /// example, every call to `read` or `write` on `TcpStream` results in a system
 282 /// call. A `BufferedStream` keeps in memory buffers of data, making large,
 283 /// infrequent calls to `read` and `write` on the underlying `Stream`.
 284 ///
 285 /// The output half will be flushed when this stream is dropped.
 286 ///
 287 /// # Example
 288 ///
 289 /// ```rust
 290 /// # #![allow(unused_must_use)]
 291 /// use std::io::{BufferedStream, File};
 292 ///
 293 /// let file = File::open(&Path::new("message.txt"));
 294 /// let mut stream = BufferedStream::new(file);
 295 ///
 296 /// stream.write("hello, world".as_bytes());
 297 /// stream.flush();
 298 ///
 299 /// let mut buf = [0, ..100];
 300 /// match stream.read(buf) {
 301 ///     Ok(nread) => println!("Read {} bytes", nread),
 302 ///     Err(e) => println!("error reading: {}", e)
 303 /// }
 304 /// ```
 305 pub struct BufferedStream<S> {
 306     inner: BufferedReader<InternalBufferedWriter<S>>
 307 }
 308
 309 impl<S: Stream> BufferedStream<S> {
 310     /// Creates a new buffered stream with explicitly listed capacities for the
 311     /// reader/writer buffer.
 312     pub fn with_capacities(reader_cap: uint, writer_cap: uint, inner: S)
 313                            -> BufferedStream<S> {
 314         let writer = BufferedWriter::with_capacity(writer_cap, inner);
 315         let internal_writer = InternalBufferedWriter(writer);
 316         let reader = BufferedReader::with_capacity(reader_cap,
 317                                                    internal_writer);
 318         BufferedStream { inner: reader }
 319     }
 320
 321     /// Creates a new buffered stream with the default reader/writer buffer
 322     /// capacities.
 323     pub fn new(inner: S) -> BufferedStream<S> {
 324         BufferedStream::with_capacities(DEFAULT_BUF_SIZE, DEFAULT_BUF_SIZE,
 325                                         inner)
 326     }
 327
 328     /// Gets a reference to the underlying stream.
 329     ///
 330     /// This type does not expose the ability to get a mutable reference to the
 331     /// underlying reader because that could possibly corrupt the buffer.
 332     pub fn get_ref<'a>(&'a self) -> &'a S {
 333         let InternalBufferedWriter(ref w) = self.inner.inner;
 334         w.get_ref()
 335     }
 336
 337     /// Unwraps this `BufferedStream`, returning the underlying stream.
 338     ///
 339     /// The internal buffer is flushed before returning the stream. Any leftover
 340     /// data in the read buffer is lost.
 341     pub fn unwrap(self) -> S {
 342         let InternalBufferedWriter(w) = self.inner.inner;
 343         w.unwrap()
 344     }
 345 }
 346
 347 impl<S: Stream> Buffer for BufferedStream<S> {
 348     fn fill_buf<'a>(&'a mut self) -> IoResult<&'a [u8]> { self.inner.fill_buf() }
 349     fn consume(&mut self, amt: uint) { self.inner.consume(amt) }
 350 }
 351
 352 impl<S: Stream> Reader for BufferedStream<S> {
 353     fn read(&mut self, buf: &mut [u8]) -> IoResult<uint> {
 354         self.inner.read(buf)
 355     }
 356 }
 357
 358 impl<S: Stream> Writer for BufferedStream<S> {
 359     fn write(&mut self, buf: &[u8]) -> IoResult<()> {
 360         self.inner.inner.get_mut_ref().write(buf)
 361     }
 362     fn flush(&mut self) -> IoResult<()> {
 363         self.inner.inner.get_mut_ref().flush()
 364     }
 365 }
 366
 367 #[cfg(test)]
 368 mod test {
 369     extern crate test;
 370     use io;
 371     use prelude::*;
 372     use super::*;
 373     use super::super::mem::{MemReader, MemWriter, BufReader};
 374     use self::test::Bencher;
 375
 376     /// A type, free to create, primarily intended for benchmarking creation of
 377     /// wrappers that, just for construction, don't need a Reader/Writer that
 378     /// does anything useful. Is equivalent to `/dev/null` in semantics.
 379     #[deriving(Clone,Eq,Ord)]
 380     pub struct NullStream;
 381
 382     impl Reader for NullStream {
 383         fn read(&mut self, _: &mut [u8]) -> io::IoResult<uint> {
 384             Err(io::standard_error(io::EndOfFile))
 385         }
 386     }
 387
 388     impl Writer for NullStream {
 389         fn write(&mut self, _: &[u8]) -> io::IoResult<()> { Ok(()) }
 390     }
 391
 392     /// A dummy reader intended at testing short-reads propagation.
 393     pub struct ShortReader {
 394         lengths: Vec<uint>,
 395     }
 396
 397     impl Reader for ShortReader {
 398         fn read(&mut self, _: &mut [u8]) -> io::IoResult<uint> {
 399             match self.lengths.shift() {
 400                 Some(i) => Ok(i),
 401                 None => Err(io::standard_error(io::EndOfFile))
 402             }
 403         }
 404     }
 405
 406     #[test]
 407     fn test_buffered_reader() {
 408         let inner = MemReader::new(vec!(0, 1, 2, 3, 4));
 409         let mut reader = BufferedReader::with_capacity(2, inner);
 410
 411         let mut buf = [0, 0, 0];
 412         let nread = reader.read(buf);
 413         assert_eq!(Ok(2), nread);
 414         assert_eq!(buf.as_slice(), &[0, 1, 0]);
 415
 416         let mut buf = [0];
 417         let nread = reader.read(buf);
 418         assert_eq!(Ok(1), nread);
 419         assert_eq!(buf.as_slice(), &[2]);
 420
 421         let mut buf = [0, 0, 0];
 422         let nread = reader.read(buf);
 423         assert_eq!(Ok(1), nread);
 424         assert_eq!(buf.as_slice(), &[3, 0, 0]);
 425
 426         let nread = reader.read(buf);
 427         assert_eq!(Ok(1), nread);
 428         assert_eq!(buf.as_slice(), &[4, 0, 0]);
 429
 430         assert!(reader.read(buf).is_err());
 431     }
 432
 433     #[test]
 434     fn test_buffered_writer() {
 435         let inner = MemWriter::new();
 436         let mut writer = BufferedWriter::with_capacity(2, inner);
 437
 438         writer.write([0, 1]).unwrap();
 439         assert_eq!(writer.get_ref().get_ref(), &[]);
 440
 441         writer.write([2]).unwrap();
 442         assert_eq!(writer.get_ref().get_ref(), &[0, 1]);
 443
 444         writer.write([3]).unwrap();
 445         assert_eq!(writer.get_ref().get_ref(), &[0, 1]);
 446
 447         writer.flush().unwrap();
 448         assert_eq!(&[0, 1, 2, 3], writer.get_ref().get_ref());
 449
 450         writer.write([4]).unwrap();
 451         writer.write([5]).unwrap();
 452         assert_eq!(&[0, 1, 2, 3], writer.get_ref().get_ref());
 453
 454         writer.write([6]).unwrap();
 455         assert_eq!(&[0, 1, 2, 3, 4, 5],
 456                    writer.get_ref().get_ref());
 457
 458         writer.write([7, 8]).unwrap();
 459         assert_eq!(&[0, 1, 2, 3, 4, 5, 6],
 460                    writer.get_ref().get_ref());
 461
 462         writer.write([9, 10, 11]).unwrap();
 463         assert_eq!(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
 464                    writer.get_ref().get_ref());
 465
 466         writer.flush().unwrap();
 467         assert_eq!(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
 468                    writer.get_ref().get_ref());
 469     }
 470
 471     #[test]
 472     fn test_buffered_writer_inner_flushes() {
 473         let mut w = BufferedWriter::with_capacity(3, MemWriter::new());
 474         w.write([0, 1]).unwrap();
 475         assert_eq!(&[], w.get_ref().get_ref());
 476         let w = w.unwrap();
 477         assert_eq!(&[0, 1], w.get_ref());
 478     }
 479
 480     // This is just here to make sure that we don't infinite loop in the
 481     // newtype struct autoderef weirdness
 482     #[test]
 483     fn test_buffered_stream() {
 484         struct S;
 485
 486         impl io::Writer for S {
 487             fn write(&mut self, _: &[u8]) -> io::IoResult<()> { Ok(()) }
 488         }
 489
 490         impl io::Reader for S {
 491             fn read(&mut self, _: &mut [u8]) -> io::IoResult<uint> {
 492                 Err(io::standard_error(io::EndOfFile))
 493             }
 494         }
 495
 496         let mut stream = BufferedStream::new(S);
 497         let mut buf = [];
 498         assert!(stream.read(buf).is_err());
 499         stream.write(buf).unwrap();
 500         stream.flush().unwrap();
 501     }
 502
 503     #[test]
 504     fn test_read_until() {
 505         let inner = MemReader::new(vec!(0, 1, 2, 1, 0));
 506         let mut reader = BufferedReader::with_capacity(2, inner);
 507         assert_eq!(reader.read_until(0), Ok(vec!(0)));
 508         assert_eq!(reader.read_until(2), Ok(vec!(1, 2)));
 509         assert_eq!(reader.read_until(1), Ok(vec!(1)));
 510         assert_eq!(reader.read_until(8), Ok(vec!(0)));
 511         assert!(reader.read_until(9).is_err());
 512     }
 513
 514     #[test]
 515     fn test_line_buffer() {
 516         let mut writer = LineBufferedWriter::new(MemWriter::new());
 517         writer.write([0]).unwrap();
 518         assert_eq!(writer.get_ref().get_ref(), &[]);
 519         writer.write([1]).unwrap();
 520         assert_eq!(writer.get_ref().get_ref(), &[]);
 521         writer.flush().unwrap();
 522         assert_eq!(writer.get_ref().get_ref(), &[0, 1]);
 523         writer.write([0, '\n' as u8, 1, '\n' as u8, 2]).unwrap();
 524         assert_eq!(writer.get_ref().get_ref(),
 525                    &[0, 1, 0, '\n' as u8, 1, '\n' as u8]);
 526         writer.flush().unwrap();
 527         assert_eq!(writer.get_ref().get_ref(),
 528                    &[0, 1, 0, '\n' as u8, 1, '\n' as u8, 2]);
 529         writer.write([3, '\n' as u8]).unwrap();
 530         assert_eq!(writer.get_ref().get_ref(),
 531             &[0, 1, 0, '\n' as u8, 1, '\n' as u8, 2, 3, '\n' as u8]);
 532     }
 533
 534     #[test]
 535     fn test_read_line() {
 536         let in_buf = MemReader::new(Vec::from_slice(bytes!("a\nb\nc")));
 537         let mut reader = BufferedReader::with_capacity(2, in_buf);
 538         assert_eq!(reader.read_line(), Ok(~"a\n"));
 539         assert_eq!(reader.read_line(), Ok(~"b\n"));
 540         assert_eq!(reader.read_line(), Ok(~"c"));
 541         assert!(reader.read_line().is_err());
 542     }
 543
 544     #[test]
 545     fn test_lines() {
 546         let in_buf = MemReader::new(Vec::from_slice(bytes!("a\nb\nc")));
 547         let mut reader = BufferedReader::with_capacity(2, in_buf);
 548         let mut it = reader.lines();
 549         assert_eq!(it.next(), Some(Ok(~"a\n")));
 550         assert_eq!(it.next(), Some(Ok(~"b\n")));
 551         assert_eq!(it.next(), Some(Ok(~"c")));
 552         assert_eq!(it.next(), None);
 553     }
 554
 555     #[test]
 556     fn test_short_reads() {
 557         let inner = ShortReader{lengths: vec![0, 1, 2, 0, 1, 0]};
 558         let mut reader = BufferedReader::new(inner);
 559         let mut buf = [0, 0];
 560         assert_eq!(reader.read(buf), Ok(0));
 561         assert_eq!(reader.read(buf), Ok(1));
 562         assert_eq!(reader.read(buf), Ok(2));
 563         assert_eq!(reader.read(buf), Ok(0));
 564         assert_eq!(reader.read(buf), Ok(1));
 565         assert_eq!(reader.read(buf), Ok(0));
 566         assert!(reader.read(buf).is_err());
 567     }
 568
 569     #[test]
 570     fn read_char_buffered() {
 571         let buf = [195u8, 159u8];
 572         let mut reader = BufferedReader::with_capacity(1, BufReader::new(buf));
 573         assert_eq!(reader.read_char(), Ok('ß'));
 574     }
 575
 576     #[test]
 577     fn test_chars() {
 578         let buf = [195u8, 159u8, 'a' as u8];
 579         let mut reader = BufferedReader::with_capacity(1, BufReader::new(buf));
 580         let mut it = reader.chars();
 581         assert_eq!(it.next(), Some(Ok('ß')));
 582         assert_eq!(it.next(), Some(Ok('a')));
 583         assert_eq!(it.next(), None);
 584     }
 585
 586     #[bench]
 587     fn bench_buffered_reader(b: &mut Bencher) {
 588         b.iter(|| {
 589             BufferedReader::new(NullStream)
 590         });
 591     }
 592
 593     #[bench]
 594     fn bench_buffered_writer(b: &mut Bencher) {
 595         b.iter(|| {
 596             BufferedWriter::new(NullStream)
 597         });
 598     }
 599
 600     #[bench]
 601     fn bench_buffered_stream(b: &mut Bencher) {
 602         b.iter(|| {
 603             BufferedStream::new(NullStream);
 604         });
 605     }
 606 }