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