src/libstd/io/util.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 //! Utility implementations of Reader and Writer
  12
  13 use prelude::v1::*;
  14 use cmp;
  15 use io;
  16 use slice::bytes::MutableByteVector;
  17
  18 /// Wraps a `Reader`, limiting the number of bytes that can be read from it.
  19 pub struct LimitReader<R> {
  20     limit: uint,
  21     inner: R
  22 }
  23
  24 impl<R: Reader> LimitReader<R> {
  25     /// Creates a new `LimitReader`
  26     pub fn new(r: R, limit: uint) -> LimitReader<R> {
  27         LimitReader { limit: limit, inner: r }
  28     }
  29
  30     /// Consumes the `LimitReader`, returning the underlying `Reader`.
  31     pub fn into_inner(self) -> R { self.inner }
  32
  33     /// Returns the number of bytes that can be read before the `LimitReader`
  34     /// will return EOF.
  35     ///
  36     /// # Note
  37     ///
  38     /// The reader may reach EOF after reading fewer bytes than indicated by
  39     /// this method if the underlying reader reaches EOF.
  40     pub fn limit(&self) -> uint { self.limit }
  41 }
  42
  43 impl<R: Reader> Reader for LimitReader<R> {
  44     fn read(&mut self, buf: &mut [u8]) -> io::IoResult<uint> {
  45         if self.limit == 0 {
  46             return Err(io::standard_error(io::EndOfFile));
  47         }
  48
  49         let len = cmp::min(self.limit, buf.len());
  50         let res = self.inner.read(buf.slice_to_mut(len));
  51         match res {
  52             Ok(len) => self.limit -= len,
  53             _ => {}
  54         }
  55         res
  56     }
  57 }
  58
  59 impl<R: Buffer> Buffer for LimitReader<R> {
  60     fn fill_buf<'a>(&'a mut self) -> io::IoResult<&'a [u8]> {
  61         let amt = try!(self.inner.fill_buf());
  62         let buf = amt[..cmp::min(amt.len(), self.limit)];
  63         if buf.len() == 0 {
  64             Err(io::standard_error(io::EndOfFile))
  65         } else {
  66             Ok(buf)
  67         }
  68     }
  69
  70     fn consume(&mut self, amt: uint) {
  71         // Don't let callers reset the limit by passing an overlarge value
  72         let amt = cmp::min(amt, self.limit);
  73         self.limit -= amt;
  74         self.inner.consume(amt);
  75     }
  76
  77 }
  78
  79 /// A `Writer` which ignores bytes written to it, like /dev/null.
  80 #[derive(Copy)]
  81 pub struct NullWriter;
  82
  83 impl Writer for NullWriter {
  84     #[inline]
  85     fn write(&mut self, _buf: &[u8]) -> io::IoResult<()> { Ok(()) }
  86 }
  87
  88 /// A `Reader` which returns an infinite stream of 0 bytes, like /dev/zero.
  89 #[derive(Copy)]
  90 pub struct ZeroReader;
  91
  92 impl Reader for ZeroReader {
  93     #[inline]
  94     fn read(&mut self, buf: &mut [u8]) -> io::IoResult<uint> {
  95         buf.set_memory(0);
  96         Ok(buf.len())
  97     }
  98 }
  99
 100 impl Buffer for ZeroReader {
 101     fn fill_buf<'a>(&'a mut self) -> io::IoResult<&'a [u8]> {
 102         static DATA: [u8; 64] = [0; 64];
 103         Ok(DATA.as_slice())
 104     }
 105
 106     fn consume(&mut self, _amt: uint) {}
 107 }
 108
 109 /// A `Reader` which is always at EOF, like /dev/null.
 110 #[derive(Copy)]
 111 pub struct NullReader;
 112
 113 impl Reader for NullReader {
 114     #[inline]
 115     fn read(&mut self, _buf: &mut [u8]) -> io::IoResult<uint> {
 116         Err(io::standard_error(io::EndOfFile))
 117     }
 118 }
 119
 120 impl Buffer for NullReader {
 121     fn fill_buf<'a>(&'a mut self) -> io::IoResult<&'a [u8]> {
 122         Err(io::standard_error(io::EndOfFile))
 123     }
 124     fn consume(&mut self, _amt: uint) {}
 125 }
 126
 127 /// A `Writer` which multiplexes writes to a set of `Writer`s.
 128 ///
 129 /// The `Writer`s are delegated to in order. If any `Writer` returns an error,
 130 /// that error is returned immediately and remaining `Writer`s are not called.
 131 pub struct MultiWriter {
 132     writers: Vec<Box<Writer+'static>>
 133 }
 134
 135 impl MultiWriter {
 136     /// Creates a new `MultiWriter`
 137     pub fn new(writers: Vec<Box<Writer+'static>>) -> MultiWriter {
 138         MultiWriter { writers: writers }
 139     }
 140 }
 141
 142 impl Writer for MultiWriter {
 143     #[inline]
 144     fn write(&mut self, buf: &[u8]) -> io::IoResult<()> {
 145         for writer in self.writers.iter_mut() {
 146             try!(writer.write(buf));
 147         }
 148         Ok(())
 149     }
 150
 151     #[inline]
 152     fn flush(&mut self) -> io::IoResult<()> {
 153         for writer in self.writers.iter_mut() {
 154             try!(writer.flush());
 155         }
 156         Ok(())
 157     }
 158 }
 159
 160 /// A `Reader` which chains input from multiple `Reader`s, reading each to
 161 /// completion before moving onto the next.
 162 #[derive(Clone)]
 163 pub struct ChainedReader<I, R> {
 164     readers: I,
 165     cur_reader: Option<R>,
 166 }
 167
 168 impl<R: Reader, I: Iterator<Item=R>> ChainedReader<I, R> {
 169     /// Creates a new `ChainedReader`
 170     pub fn new(mut readers: I) -> ChainedReader<I, R> {
 171         let r = readers.next();
 172         ChainedReader { readers: readers, cur_reader: r }
 173     }
 174 }
 175
 176 impl<R: Reader, I: Iterator<Item=R>> Reader for ChainedReader<I, R> {
 177     fn read(&mut self, buf: &mut [u8]) -> io::IoResult<uint> {
 178         loop {
 179             let err = match self.cur_reader {
 180                 Some(ref mut r) => {
 181                     match r.read(buf) {
 182                         Ok(len) => return Ok(len),
 183                         Err(ref e) if e.kind == io::EndOfFile => None,
 184                         Err(e) => Some(e),
 185                     }
 186                 }
 187                 None => break
 188             };
 189             self.cur_reader = self.readers.next();
 190             match err {
 191                 Some(e) => return Err(e),
 192                 None => {}
 193             }
 194         }
 195         Err(io::standard_error(io::EndOfFile))
 196     }
 197 }
 198
 199 /// A `Reader` which forwards input from another `Reader`, passing it along to
 200 /// a `Writer` as well. Similar to the `tee(1)` command.
 201 pub struct TeeReader<R, W> {
 202     reader: R,
 203     writer: W,
 204 }
 205
 206 impl<R: Reader, W: Writer> TeeReader<R, W> {
 207     /// Creates a new `TeeReader`
 208     pub fn new(r: R, w: W) -> TeeReader<R, W> {
 209         TeeReader { reader: r, writer: w }
 210     }
 211
 212     /// Consumes the `TeeReader`, returning the underlying `Reader` and
 213     /// `Writer`.
 214     pub fn into_inner(self) -> (R, W) {
 215         let TeeReader { reader, writer } = self;
 216         (reader, writer)
 217     }
 218 }
 219
 220 impl<R: Reader, W: Writer> Reader for TeeReader<R, W> {
 221     fn read(&mut self, buf: &mut [u8]) -> io::IoResult<uint> {
 222         self.reader.read(buf).and_then(|len| {
 223             self.writer.write(buf[mut ..len]).map(|()| len)
 224         })
 225     }
 226 }
 227
 228 /// Copies all data from a `Reader` to a `Writer`.
 229 pub fn copy<R: Reader, W: Writer>(r: &mut R, w: &mut W) -> io::IoResult<()> {
 230     let mut buf = [0; super::DEFAULT_BUF_SIZE];
 231     loop {
 232         let len = match r.read(&mut buf) {
 233             Ok(len) => len,
 234             Err(ref e) if e.kind == io::EndOfFile => return Ok(()),
 235             Err(e) => return Err(e),
 236         };
 237         try!(w.write(buf[..len]));
 238     }
 239 }
 240
 241 /// An adaptor converting an `Iterator<u8>` to a `Reader`.
 242 #[derive(Clone)]
 243 pub struct IterReader<T> {
 244     iter: T,
 245 }
 246
 247 impl<T: Iterator<Item=u8>> IterReader<T> {
 248     /// Creates a new `IterReader` which will read from the specified
 249     /// `Iterator`.
 250     pub fn new(iter: T) -> IterReader<T> {
 251         IterReader { iter: iter }
 252     }
 253 }
 254
 255 impl<T: Iterator<Item=u8>> Reader for IterReader<T> {
 256     #[inline]
 257     fn read(&mut self, buf: &mut [u8]) -> io::IoResult<uint> {
 258         let mut len = 0;
 259         for (slot, elt) in buf.iter_mut().zip(self.iter.by_ref()) {
 260             *slot = elt;
 261             len += 1;
 262         }
 263         if len == 0 && buf.len() != 0 {
 264             Err(io::standard_error(io::EndOfFile))
 265         } else {
 266             Ok(len)
 267         }
 268     }
 269 }
 270
 271 #[cfg(test)]
 272 mod test {
 273     use prelude::v1::*;
 274
 275     use io::{MemReader, ByRefReader};
 276     use io;
 277     use super::*;
 278
 279     #[test]
 280     fn test_limit_reader_unlimited() {
 281         let mut r = MemReader::new(vec!(0, 1, 2));
 282         {
 283             let mut r = LimitReader::new(r.by_ref(), 4);
 284             assert_eq!(vec!(0, 1, 2), r.read_to_end().unwrap());
 285         }
 286     }
 287
 288     #[test]
 289     fn test_limit_reader_limited() {
 290         let mut r = MemReader::new(vec!(0, 1, 2));
 291         {
 292             let mut r = LimitReader::new(r.by_ref(), 2);
 293             assert_eq!(vec!(0, 1), r.read_to_end().unwrap());
 294         }
 295         assert_eq!(vec!(2), r.read_to_end().unwrap());
 296     }
 297
 298     #[test]
 299     fn test_limit_reader_limit() {
 300         let r = MemReader::new(vec!(0, 1, 2));
 301         let mut r = LimitReader::new(r, 3);
 302         assert_eq!(3, r.limit());
 303         assert_eq!(0, r.read_byte().unwrap());
 304         assert_eq!(2, r.limit());
 305         assert_eq!(vec!(1, 2), r.read_to_end().unwrap());
 306         assert_eq!(0, r.limit());
 307     }
 308
 309     #[test]
 310     fn test_limit_reader_overlong_consume() {
 311         let mut r = MemReader::new(vec![0, 1, 2, 3, 4, 5]);
 312         let mut r = LimitReader::new(r.by_ref(), 1);
 313         r.consume(2);
 314         assert_eq!(vec![], r.read_to_end().unwrap());
 315     }
 316
 317     #[test]
 318     fn test_null_writer() {
 319         let mut s = NullWriter;
 320         let buf = vec![0, 0, 0];
 321         s.write(buf.as_slice()).unwrap();
 322         s.flush().unwrap();
 323     }
 324
 325     #[test]
 326     fn test_zero_reader() {
 327         let mut s = ZeroReader;
 328         let mut buf = vec![1, 2, 3];
 329         assert_eq!(s.read(buf.as_mut_slice()), Ok(3));
 330         assert_eq!(vec![0, 0, 0], buf);
 331     }
 332
 333     #[test]
 334     fn test_null_reader() {
 335         let mut r = NullReader;
 336         let mut buf = vec![0];
 337         assert!(r.read(buf.as_mut_slice()).is_err());
 338     }
 339
 340     #[test]
 341     fn test_multi_writer() {
 342         static mut writes: uint = 0;
 343         static mut flushes: uint = 0;
 344
 345         struct TestWriter;
 346         impl Writer for TestWriter {
 347             fn write(&mut self, _buf: &[u8]) -> io::IoResult<()> {
 348                 unsafe { writes += 1 }
 349                 Ok(())
 350             }
 351
 352             fn flush(&mut self) -> io::IoResult<()> {
 353                 unsafe { flushes += 1 }
 354                 Ok(())
 355             }
 356         }
 357
 358         let mut multi = MultiWriter::new(vec!(box TestWriter as Box<Writer>,
 359                                               box TestWriter as Box<Writer>));
 360         multi.write(&[1, 2, 3]).unwrap();
 361         assert_eq!(2, unsafe { writes });
 362         assert_eq!(0, unsafe { flushes });
 363         multi.flush().unwrap();
 364         assert_eq!(2, unsafe { writes });
 365         assert_eq!(2, unsafe { flushes });
 366     }
 367
 368     #[test]
 369     fn test_chained_reader() {
 370         let rs = vec!(MemReader::new(vec!(0, 1)), MemReader::new(vec!()),
 371                       MemReader::new(vec!(2, 3)));
 372         let mut r = ChainedReader::new(rs.into_iter());
 373         assert_eq!(vec!(0, 1, 2, 3), r.read_to_end().unwrap());
 374     }
 375
 376     #[test]
 377     fn test_tee_reader() {
 378         let mut r = TeeReader::new(MemReader::new(vec!(0, 1, 2)),
 379                                    Vec::new());
 380         assert_eq!(vec!(0, 1, 2), r.read_to_end().unwrap());
 381         let (_, w) = r.into_inner();
 382         assert_eq!(vec!(0, 1, 2), w);
 383     }
 384
 385     #[test]
 386     fn test_copy() {
 387         let mut r = MemReader::new(vec!(0, 1, 2, 3, 4));
 388         let mut w = Vec::new();
 389         copy(&mut r, &mut w).unwrap();
 390         assert_eq!(vec!(0, 1, 2, 3, 4), w);
 391     }
 392
 393     #[test]
 394     fn limit_reader_buffer() {
 395         let r = &mut b"0123456789\n0123456789\n";
 396         {
 397             let mut r = LimitReader::new(r.by_ref(), 3);
 398             assert_eq!(r.read_line(), Ok("012".to_string()));
 399             assert_eq!(r.limit(), 0);
 400             assert_eq!(r.read_line().err().unwrap().kind, io::EndOfFile);
 401         }
 402         {
 403             let mut r = LimitReader::new(r.by_ref(), 9);
 404             assert_eq!(r.read_line(), Ok("3456789\n".to_string()));
 405             assert_eq!(r.limit(), 1);
 406             assert_eq!(r.read_line(), Ok("0".to_string()));
 407         }
 408         {
 409             let mut r = LimitReader::new(r.by_ref(), 100);
 410             assert_eq!(r.read_char(), Ok('1'));
 411             assert_eq!(r.limit(), 99);
 412             assert_eq!(r.read_line(), Ok("23456789\n".to_string()));
 413         }
 414     }
 415
 416     #[test]
 417     fn test_iter_reader() {
 418         let mut r = IterReader::new(range(0u8, 8));
 419         let mut buf = [0, 0, 0];
 420         let len = r.read(&mut buf).unwrap();
 421         assert_eq!(len, 3);
 422         assert!(buf == [0, 1, 2]);
 423
 424         let len = r.read(&mut buf).unwrap();
 425         assert_eq!(len, 3);
 426         assert!(buf == [3, 4, 5]);
 427
 428         let len = r.read(&mut buf).unwrap();
 429         assert_eq!(len, 2);
 430         assert!(buf == [6, 7, 5]);
 431
 432         assert_eq!(r.read(&mut buf).unwrap_err().kind, io::EndOfFile);
 433     }
 434
 435     #[test]
 436     fn iter_reader_zero_length() {
 437         let mut r = IterReader::new(range(0u8, 8));
 438         let mut buf = [];
 439         assert_eq!(Ok(0), r.read(&mut buf));
 440     }
 441 }