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