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