From 7fe5e3082394ba73f54ceeadd3ed092fd54ee0ff Mon Sep 17 00:00:00 2001 From: Steven Fackler Date: Mon, 9 Dec 2013 22:53:09 -0800 Subject: [PATCH] Add std::io::util This adds a bunch of useful Reader and Writer implementations. I'm not a huge fan of the name `util` but I can't think of a better name and I don't want to make `std::io` any longer than it already is. --- src/libstd/io/mod.rs | 3 + src/libstd/io/util.rs | 295 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 298 insertions(+) create mode 100644 src/libstd/io/util.rs diff --git a/src/libstd/io/mod.rs b/src/libstd/io/mod.rs index edc6728a0c1d7..56193bf2754e7 100644 --- a/src/libstd/io/mod.rs +++ b/src/libstd/io/mod.rs @@ -312,6 +312,9 @@ pub mod native; /// Signal handling pub mod signal; +/// Utility implementations of Reader and Writer +pub mod util; + /// The default buffer size for various I/O operations static DEFAULT_BUF_SIZE: uint = 1024 * 64; diff --git a/src/libstd/io/util.rs b/src/libstd/io/util.rs new file mode 100644 index 0000000000000..75be56113e930 --- /dev/null +++ b/src/libstd/io/util.rs @@ -0,0 +1,295 @@ +// Copyright 2013 The Rust Project Developers. See the COPYRIGHT +// file at the top-level directory of this distribution and at +// http://rust-lang.org/COPYRIGHT. +// +// Licensed under the Apache License, Version 2.0 or the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. +use prelude::*; +use cmp; +use vec::bytes::MutableByteVector; + +/// Wraps a `Reader`, limiting the number of bytes that can be read from it. +pub struct LimitReader<'a, R> { + priv limit: uint, + priv inner: &'a mut R +} + +impl<'a, R: Reader> LimitReader<'a, R> { + /// Creates a new `LimitReader` + pub fn new<'a>(r: &'a mut R, limit: uint) -> LimitReader<'a, R> { + LimitReader { limit: limit, inner: r } + } +} + +impl<'a, R: Reader> Reader for LimitReader<'a, R> { + fn read(&mut self, buf: &mut [u8]) -> Option { + if self.limit == 0 { + return None; + } + + let len = cmp::min(self.limit, buf.len()); + self.inner.read(buf.mut_slice_to(len)).map(|len| { + self.limit -= len; + len + }) + } + + fn eof(&mut self) -> bool { + self.limit == 0 || self.inner.eof() + } +} + +/// A `Writer` which ignores bytes written to it, like /dev/null. +pub struct NullWriter; + +impl Writer for NullWriter { + #[inline] + fn write(&mut self, _buf: &[u8]) { } +} + +/// A `Reader` which returns an infinite stream of 0 bytes, like /dev/zero. +pub struct ZeroReader; + +impl Reader for ZeroReader { + #[inline] + fn read(&mut self, buf: &mut [u8]) -> Option { + buf.set_memory(0); + Some(buf.len()) + } + + #[inline] + fn eof(&mut self) -> bool { + false + } +} + +/// A `Reader` which is always at EOF, like /dev/null. +pub struct NullReader; + +impl Reader for NullReader { + #[inline] + fn read(&mut self, _buf: &mut [u8]) -> Option { + None + } + + #[inline] + fn eof(&mut self) -> bool { + true + } +} + +/// A `Writer` which multiplexes writes to a set of `Writers`. +pub struct MultiWriter { + priv writers: ~[~Writer] +} + +impl MultiWriter { + /// Creates a new `MultiWriter` + pub fn new(writers: ~[~Writer]) -> MultiWriter { + MultiWriter { writers: writers } + } +} + +impl Writer for MultiWriter { + #[inline] + fn write(&mut self, buf: &[u8]) { + for writer in self.writers.mut_iter() { + writer.write(buf); + } + } + + #[inline] + fn flush(&mut self) { + for writer in self.writers.mut_iter() { + writer.flush(); + } + } +} + +/// A `Reader` which chains input from multiple `Readers`, reading each to +/// completion before moving onto the next. +pub struct ChainedReader { + priv readers: I, + priv cur_reader: Option, +} + +impl> ChainedReader { + /// Creates a new `ChainedReader` + pub fn new(mut readers: I) -> ChainedReader { + let r = readers.next(); + ChainedReader { readers: readers, cur_reader: r } + } +} + +impl> Reader for ChainedReader { + fn read(&mut self, buf: &mut [u8]) -> Option { + loop { + match self.cur_reader { + Some(ref mut r) => { + match r.read(buf) { + Some(len) => return Some(len), + None => {} + } + } + None => break + } + self.cur_reader = self.readers.next() + } + None + } + + fn eof(&mut self) -> bool { + self.cur_reader.is_none() + } +} + +/// A `Reader` which forwards input from another `Reader`, passing it along to +/// a `Writer` as well. Similar to the `tee(1)` command. +pub struct TeeReader { + priv reader: R, + priv writer: W +} + +impl TeeReader { + /// Creates a new `TeeReader` + pub fn new(r: R, w: W) -> TeeReader { + TeeReader { reader: r, writer: w } + } + + /// Consumes the `TeeReader`, returning the underlying `Reader` and + /// `Writer`. + pub fn unwrap(self) -> (R, W) { + let TeeReader { reader, writer } = self; + (reader, writer) + } +} + +impl Reader for TeeReader { + fn read(&mut self, buf: &mut [u8]) -> Option { + self.reader.read(buf).map(|len| { + self.writer.write(buf.slice_to(len)); + len + }) + } + + fn eof(&mut self) -> bool { + self.reader.eof() + } +} + +/// Copies all data from a `Reader` to a `Writer`. +pub fn copy(r: &mut R, w: &mut W) { + let mut buf = [0, ..super::DEFAULT_BUF_SIZE]; + loop { + match r.read(buf) { + Some(len) => w.write(buf.slice_to(len)), + None => break + } + } +} + +#[cfg(test)] +mod test { + use io::mem::{MemReader, MemWriter}; + use super::*; + use prelude::*; + + #[test] + fn test_bounded_reader_unlimited() { + let mut r = MemReader::new(~[0, 1, 2]); + { + let mut r = LimitReader::new(&mut r, 4); + assert_eq!(~[0, 1, 2], r.read_to_end()); + } + assert!(r.eof()); + } + + #[test] + fn test_bound_reader_limited() { + let mut r = MemReader::new(~[0, 1, 2]); + { + let mut r = LimitReader::new(&mut r, 2); + assert_eq!(~[0, 1], r.read_to_end()); + } + assert_eq!(~[2], r.read_to_end()); + } + + #[test] + fn test_null_writer() { + let mut s = NullWriter; + let mut buf = ~[0, 0, 0]; + s.write(buf); + s.flush(); + } + + #[test] + fn test_zero_reader() { + let mut s = ZeroReader; + let mut buf = ~[1, 2, 3]; + assert_eq!(s.read(buf), Some(3)); + assert_eq!(~[0, 0, 0], buf); + } + + #[test] + fn test_null_reader() { + let mut r = NullReader; + let mut buf = ~[0]; + assert_eq!(r.read(buf), None); + assert!(r.eof()); + } + + #[test] + fn test_multi_writer() { + static mut writes: uint = 0; + static mut flushes: uint = 0; + + struct TestWriter; + impl Writer for TestWriter { + fn write(&mut self, buf: &[u8]) { + unsafe { writes += 1 } + } + + fn flush(&mut self) { + unsafe { flushes += 1 } + } + } + + let mut multi = MultiWriter::new(~[~TestWriter as ~Writer, + ~TestWriter as ~Writer]); + multi.write([1, 2, 3]); + assert_eq!(2, unsafe { writes }); + assert_eq!(0, unsafe { flushes }); + multi.flush(); + assert_eq!(2, unsafe { writes }); + assert_eq!(2, unsafe { flushes }); + } + + #[test] + fn test_chained_reader() { + let rs = ~[MemReader::new(~[0, 1]), MemReader::new(~[]), + MemReader::new(~[2, 3])]; + let mut r = ChainedReader::new(rs.move_iter()); + assert_eq!(~[0, 1, 2, 3], r.read_to_end()); + } + + #[test] + fn test_tee_reader() { + let mut r = TeeReader::new(MemReader::new(~[0, 1, 2]), + MemWriter::new()); + assert_eq!(~[0, 1, 2], r.read_to_end()); + let (_, w) = r.unwrap(); + assert_eq!(~[0, 1, 2], w.inner()); + } + + #[test] + fn test_copy() { + let mut r = MemReader::new(~[0, 1, 2, 3, 4]); + let mut w = MemWriter::new(); + copy(&mut r, &mut w); + assert!(r.eof()); + assert_eq!(~[0, 1, 2, 3, 4], w.inner()); + } +}