|
| 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 | +use prelude::*; |
| 11 | +use cmp; |
| 12 | +use vec::bytes::MutableByteVector; |
| 13 | + |
| 14 | +/// Wraps a `Reader`, limiting the number of bytes that can be read from it. |
| 15 | +pub struct LimitReader<'a, R> { |
| 16 | + priv limit: uint, |
| 17 | + priv inner: &'a mut R |
| 18 | +} |
| 19 | + |
| 20 | +impl<'a, R: Reader> LimitReader<'a, R> { |
| 21 | + /// Creates a new `LimitReader` |
| 22 | + pub fn new<'a>(r: &'a mut R, limit: uint) -> LimitReader<'a, R> { |
| 23 | + LimitReader { limit: limit, inner: r } |
| 24 | + } |
| 25 | +} |
| 26 | + |
| 27 | +impl<'a, R: Reader> Reader for LimitReader<'a, R> { |
| 28 | + fn read(&mut self, buf: &mut [u8]) -> Option<uint> { |
| 29 | + if self.limit == 0 { |
| 30 | + return None; |
| 31 | + } |
| 32 | + |
| 33 | + let len = cmp::min(self.limit, buf.len()); |
| 34 | + self.inner.read(buf.mut_slice_to(len)).map(|len| { |
| 35 | + self.limit -= len; |
| 36 | + len |
| 37 | + }) |
| 38 | + } |
| 39 | + |
| 40 | + fn eof(&mut self) -> bool { |
| 41 | + self.limit == 0 || self.inner.eof() |
| 42 | + } |
| 43 | +} |
| 44 | + |
| 45 | +/// A `Writer` which ignores bytes written to it, like /dev/null. |
| 46 | +pub struct NullWriter; |
| 47 | + |
| 48 | +impl Writer for NullWriter { |
| 49 | + #[inline] |
| 50 | + fn write(&mut self, _buf: &[u8]) { } |
| 51 | +} |
| 52 | + |
| 53 | +/// A `Reader` which returns an infinite stream of 0 bytes, like /dev/zero. |
| 54 | +pub struct ZeroReader; |
| 55 | + |
| 56 | +impl Reader for ZeroReader { |
| 57 | + #[inline] |
| 58 | + fn read(&mut self, buf: &mut [u8]) -> Option<uint> { |
| 59 | + buf.set_memory(0); |
| 60 | + Some(buf.len()) |
| 61 | + } |
| 62 | + |
| 63 | + #[inline] |
| 64 | + fn eof(&mut self) -> bool { |
| 65 | + false |
| 66 | + } |
| 67 | +} |
| 68 | + |
| 69 | +/// A `Reader` which is always at EOF, like /dev/null. |
| 70 | +pub struct NullReader; |
| 71 | + |
| 72 | +impl Reader for NullReader { |
| 73 | + #[inline] |
| 74 | + fn read(&mut self, _buf: &mut [u8]) -> Option<uint> { |
| 75 | + None |
| 76 | + } |
| 77 | + |
| 78 | + #[inline] |
| 79 | + fn eof(&mut self) -> bool { |
| 80 | + true |
| 81 | + } |
| 82 | +} |
| 83 | + |
| 84 | +/// A `Writer` which multiplexes writes to a set of `Writers`. |
| 85 | +pub struct MultiWriter { |
| 86 | + priv writers: ~[~Writer] |
| 87 | +} |
| 88 | + |
| 89 | +impl MultiWriter { |
| 90 | + /// Creates a new `MultiWriter` |
| 91 | + pub fn new(writers: ~[~Writer]) -> MultiWriter { |
| 92 | + MultiWriter { writers: writers } |
| 93 | + } |
| 94 | +} |
| 95 | + |
| 96 | +impl Writer for MultiWriter { |
| 97 | + #[inline] |
| 98 | + fn write(&mut self, buf: &[u8]) { |
| 99 | + for writer in self.writers.mut_iter() { |
| 100 | + writer.write(buf); |
| 101 | + } |
| 102 | + } |
| 103 | + |
| 104 | + #[inline] |
| 105 | + fn flush(&mut self) { |
| 106 | + for writer in self.writers.mut_iter() { |
| 107 | + writer.flush(); |
| 108 | + } |
| 109 | + } |
| 110 | +} |
| 111 | + |
| 112 | +/// A `Reader` which chains input from multiple `Readers`, reading each to |
| 113 | +/// completion before moving onto the next. |
| 114 | +pub struct ChainedReader<I, R> { |
| 115 | + priv readers: I, |
| 116 | + priv cur_reader: Option<R>, |
| 117 | +} |
| 118 | + |
| 119 | +impl<R: Reader, I: Iterator<R>> ChainedReader<I, R> { |
| 120 | + /// Creates a new `ChainedReader` |
| 121 | + pub fn new(mut readers: I) -> ChainedReader<I, R> { |
| 122 | + let r = readers.next(); |
| 123 | + ChainedReader { readers: readers, cur_reader: r } |
| 124 | + } |
| 125 | +} |
| 126 | + |
| 127 | +impl<R: Reader, I: Iterator<R>> Reader for ChainedReader<I, R> { |
| 128 | + fn read(&mut self, buf: &mut [u8]) -> Option<uint> { |
| 129 | + loop { |
| 130 | + match self.cur_reader { |
| 131 | + Some(ref mut r) => { |
| 132 | + match r.read(buf) { |
| 133 | + Some(len) => return Some(len), |
| 134 | + None => {} |
| 135 | + } |
| 136 | + } |
| 137 | + None => break |
| 138 | + } |
| 139 | + self.cur_reader = self.readers.next() |
| 140 | + } |
| 141 | + None |
| 142 | + } |
| 143 | + |
| 144 | + fn eof(&mut self) -> bool { |
| 145 | + self.cur_reader.is_none() |
| 146 | + } |
| 147 | +} |
| 148 | + |
| 149 | +/// A `Reader` which forwards input from another `Reader`, passing it along to |
| 150 | +/// a `Writer` as well. Similar to the `tee(1)` command. |
| 151 | +pub struct TeeReader<R, W> { |
| 152 | + priv reader: R, |
| 153 | + priv writer: W |
| 154 | +} |
| 155 | + |
| 156 | +impl<R: Reader, W: Writer> TeeReader<R, W> { |
| 157 | + /// Creates a new `TeeReader` |
| 158 | + pub fn new(r: R, w: W) -> TeeReader<R, W> { |
| 159 | + TeeReader { reader: r, writer: w } |
| 160 | + } |
| 161 | + |
| 162 | + /// Consumes the `TeeReader`, returning the underlying `Reader` and |
| 163 | + /// `Writer`. |
| 164 | + pub fn unwrap(self) -> (R, W) { |
| 165 | + let TeeReader { reader, writer } = self; |
| 166 | + (reader, writer) |
| 167 | + } |
| 168 | +} |
| 169 | + |
| 170 | +impl<R: Reader, W: Writer> Reader for TeeReader<R, W> { |
| 171 | + fn read(&mut self, buf: &mut [u8]) -> Option<uint> { |
| 172 | + self.reader.read(buf).map(|len| { |
| 173 | + self.writer.write(buf.slice_to(len)); |
| 174 | + len |
| 175 | + }) |
| 176 | + } |
| 177 | + |
| 178 | + fn eof(&mut self) -> bool { |
| 179 | + self.reader.eof() |
| 180 | + } |
| 181 | +} |
| 182 | + |
| 183 | +/// Copies all data from a `Reader` to a `Writer`. |
| 184 | +pub fn copy<R: Reader, W: Writer>(r: &mut R, w: &mut W) { |
| 185 | + let mut buf = [0, ..super::DEFAULT_BUF_SIZE]; |
| 186 | + loop { |
| 187 | + match r.read(buf) { |
| 188 | + Some(len) => w.write(buf.slice_to(len)), |
| 189 | + None => break |
| 190 | + } |
| 191 | + } |
| 192 | +} |
| 193 | + |
| 194 | +#[cfg(test)] |
| 195 | +mod test { |
| 196 | + use io::mem::{MemReader, MemWriter}; |
| 197 | + use super::*; |
| 198 | + use prelude::*; |
| 199 | + |
| 200 | + #[test] |
| 201 | + fn test_bounded_reader_unlimited() { |
| 202 | + let mut r = MemReader::new(~[0, 1, 2]); |
| 203 | + { |
| 204 | + let mut r = LimitReader::new(&mut r, 4); |
| 205 | + assert_eq!(~[0, 1, 2], r.read_to_end()); |
| 206 | + } |
| 207 | + assert!(r.eof()); |
| 208 | + } |
| 209 | + |
| 210 | + #[test] |
| 211 | + fn test_bound_reader_limited() { |
| 212 | + let mut r = MemReader::new(~[0, 1, 2]); |
| 213 | + { |
| 214 | + let mut r = LimitReader::new(&mut r, 2); |
| 215 | + assert_eq!(~[0, 1], r.read_to_end()); |
| 216 | + } |
| 217 | + assert_eq!(~[2], r.read_to_end()); |
| 218 | + } |
| 219 | + |
| 220 | + #[test] |
| 221 | + fn test_null_writer() { |
| 222 | + let mut s = NullWriter; |
| 223 | + let mut buf = ~[0, 0, 0]; |
| 224 | + s.write(buf); |
| 225 | + s.flush(); |
| 226 | + } |
| 227 | + |
| 228 | + #[test] |
| 229 | + fn test_zero_reader() { |
| 230 | + let mut s = ZeroReader; |
| 231 | + let mut buf = ~[1, 2, 3]; |
| 232 | + assert_eq!(s.read(buf), Some(3)); |
| 233 | + assert_eq!(~[0, 0, 0], buf); |
| 234 | + } |
| 235 | + |
| 236 | + #[test] |
| 237 | + fn test_null_reader() { |
| 238 | + let mut r = NullReader; |
| 239 | + let mut buf = ~[0]; |
| 240 | + assert_eq!(r.read(buf), None); |
| 241 | + assert!(r.eof()); |
| 242 | + } |
| 243 | + |
| 244 | + #[test] |
| 245 | + fn test_multi_writer() { |
| 246 | + static mut writes: uint = 0; |
| 247 | + static mut flushes: uint = 0; |
| 248 | + |
| 249 | + struct TestWriter; |
| 250 | + impl Writer for TestWriter { |
| 251 | + fn write(&mut self, buf: &[u8]) { |
| 252 | + unsafe { writes += 1 } |
| 253 | + } |
| 254 | + |
| 255 | + fn flush(&mut self) { |
| 256 | + unsafe { flushes += 1 } |
| 257 | + } |
| 258 | + } |
| 259 | + |
| 260 | + let mut multi = MultiWriter::new(~[~TestWriter as ~Writer, |
| 261 | + ~TestWriter as ~Writer]); |
| 262 | + multi.write([1, 2, 3]); |
| 263 | + assert_eq!(2, unsafe { writes }); |
| 264 | + assert_eq!(0, unsafe { flushes }); |
| 265 | + multi.flush(); |
| 266 | + assert_eq!(2, unsafe { writes }); |
| 267 | + assert_eq!(2, unsafe { flushes }); |
| 268 | + } |
| 269 | + |
| 270 | + #[test] |
| 271 | + fn test_chained_reader() { |
| 272 | + let rs = ~[MemReader::new(~[0, 1]), MemReader::new(~[]), |
| 273 | + MemReader::new(~[2, 3])]; |
| 274 | + let mut r = ChainedReader::new(rs.move_iter()); |
| 275 | + assert_eq!(~[0, 1, 2, 3], r.read_to_end()); |
| 276 | + } |
| 277 | + |
| 278 | + #[test] |
| 279 | + fn test_tee_reader() { |
| 280 | + let mut r = TeeReader::new(MemReader::new(~[0, 1, 2]), |
| 281 | + MemWriter::new()); |
| 282 | + assert_eq!(~[0, 1, 2], r.read_to_end()); |
| 283 | + let (_, w) = r.unwrap(); |
| 284 | + assert_eq!(~[0, 1, 2], w.inner()); |
| 285 | + } |
| 286 | + |
| 287 | + #[test] |
| 288 | + fn test_copy() { |
| 289 | + let mut r = MemReader::new(~[0, 1, 2, 3, 4]); |
| 290 | + let mut w = MemWriter::new(); |
| 291 | + copy(&mut r, &mut w); |
| 292 | + assert!(r.eof()); |
| 293 | + assert_eq!(~[0, 1, 2, 3, 4], w.inner()); |
| 294 | + } |
| 295 | +} |
0 commit comments