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