src/libstd/io/mem.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 //! Readers and Writers for in-memory buffers
  12 //!
  13 //! # XXX
  14 //!
  15 //! * Should probably have something like this for strings.
  16 //! * Should they implement Closable? Would take extra state.
  17 use cmp::max;
  18 use cmp::min;
  19 use prelude::*;
  20 use super::*;
  21 use vec;
  22
  23 /// Writes to an owned, growable byte vector
  24 pub struct MemWriter {
  25     priv buf: ~[u8],
  26     priv pos: uint,
  27 }
  28
  29 impl MemWriter {
  30     /// Create a new `MemWriter`.
  31     pub fn new() -> MemWriter {
  32         MemWriter::with_capacity(128)
  33     }
  34     /// Create a new `MemWriter`, allocating at least `n` bytes for
  35     /// the internal buffer.
  36     pub fn with_capacity(n: uint) -> MemWriter {
  37         MemWriter { buf: vec::with_capacity(n), pos: 0 }
  38     }
  39
  40     /// Acquires an immutable reference to the underlying buffer of this
  41     /// `MemWriter`.
  42     ///
  43     /// No method is exposed for acquiring a mutable reference to the buffer
  44     /// because it could corrupt the state of this `MemWriter`.
  45     pub fn get_ref<'a>(&'a self) -> &'a [u8] { self.buf.as_slice() }
  46
  47     /// Unwraps this `MemWriter`, returning the underlying buffer
  48     pub fn unwrap(self) -> ~[u8] { self.buf }
  49 }
  50
  51 impl Writer for MemWriter {
  52     fn write(&mut self, buf: &[u8]) {
  53         // Make sure the internal buffer is as least as big as where we
  54         // currently are
  55         let difference = self.pos as i64 - self.buf.len() as i64;
  56         if difference > 0 {
  57             self.buf.grow(difference as uint, &0);
  58         }
  59
  60         // Figure out what bytes will be used to overwrite what's currently
  61         // there (left), and what will be appended on the end (right)
  62         let cap = self.buf.len() - self.pos;
  63         let (left, right) = if cap <= buf.len() {
  64             (buf.slice_to(cap), buf.slice_from(cap))
  65         } else {
  66             (buf, &[])
  67         };
  68
  69         // Do the necessary writes
  70         if left.len() > 0 {
  71             vec::bytes::copy_memory(self.buf.mut_slice_from(self.pos), left);
  72         }
  73         if right.len() > 0 {
  74             self.buf.push_all(right);
  75         }
  76
  77         // Bump us forward
  78         self.pos += buf.len();
  79     }
  80 }
  81
  82 // FIXME(#10432)
  83 impl Seek for MemWriter {
  84     fn tell(&self) -> u64 { self.pos as u64 }
  85
  86     fn seek(&mut self, pos: i64, style: SeekStyle) {
  87         // compute offset as signed and clamp to prevent overflow
  88         let offset = match style {
  89             SeekSet => { 0 }
  90             SeekEnd => { self.buf.len() }
  91             SeekCur => { self.pos }
  92         } as i64;
  93
  94         self.pos = max(0, offset+pos) as uint;
  95     }
  96 }
  97
  98 /// Reads from an owned byte vector
  99 pub struct MemReader {
 100     priv buf: ~[u8],
 101     priv pos: uint
 102 }
 103
 104 impl MemReader {
 105     /// Creates a new `MemReader` which will read the buffer given. The buffer
 106     /// can be re-acquired through `unwrap`
 107     pub fn new(buf: ~[u8]) -> MemReader {
 108         MemReader {
 109             buf: buf,
 110             pos: 0
 111         }
 112     }
 113
 114     /// Tests whether this reader has read all bytes in its buffer.
 115     ///
 116     /// If `true`, then this will no longer return bytes from `read`.
 117     pub fn eof(&self) -> bool { self.pos == self.buf.len() }
 118
 119     /// Acquires an immutable reference to the underlying buffer of this
 120     /// `MemReader`.
 121     ///
 122     /// No method is exposed for acquiring a mutable reference to the buffer
 123     /// because it could corrupt the state of this `MemReader`.
 124     pub fn get_ref<'a>(&'a self) -> &'a [u8] { self.buf.as_slice() }
 125
 126     /// Unwraps this `MemReader`, returning the underlying buffer
 127     pub fn unwrap(self) -> ~[u8] { self.buf }
 128 }
 129
 130 impl Reader for MemReader {
 131     fn read(&mut self, buf: &mut [u8]) -> Option<uint> {
 132         if self.eof() { return None }
 133
 134         let write_len = min(buf.len(), self.buf.len() - self.pos);
 135         {
 136             let input = self.buf.slice(self.pos, self.pos + write_len);
 137             let output = buf.mut_slice(0, write_len);
 138             assert_eq!(input.len(), output.len());
 139             vec::bytes::copy_memory(output, input);
 140         }
 141         self.pos += write_len;
 142         assert!(self.pos <= self.buf.len());
 143
 144         return Some(write_len);
 145     }
 146 }
 147
 148 impl Seek for MemReader {
 149     fn tell(&self) -> u64 { self.pos as u64 }
 150     fn seek(&mut self, _pos: i64, _style: SeekStyle) { fail!() }
 151 }
 152
 153 impl Buffer for MemReader {
 154     fn fill<'a>(&'a mut self) -> &'a [u8] { self.buf.slice_from(self.pos) }
 155     fn consume(&mut self, amt: uint) { self.pos += amt; }
 156 }
 157
 158 /// Writes to a fixed-size byte slice
 159 ///
 160 /// If a write will not fit in the buffer, it raises the `io_error`
 161 /// condition and does not write any data.
 162 pub struct BufWriter<'a> {
 163     priv buf: &'a mut [u8],
 164     priv pos: uint
 165 }
 166
 167 impl<'a> BufWriter<'a> {
 168     pub fn new<'a>(buf: &'a mut [u8]) -> BufWriter<'a> {
 169         BufWriter {
 170             buf: buf,
 171             pos: 0
 172         }
 173     }
 174 }
 175
 176 impl<'a> Writer for BufWriter<'a> {
 177     fn write(&mut self, buf: &[u8]) {
 178         // raises a condition if the entire write does not fit in the buffer
 179         let max_size = self.buf.len();
 180         if self.pos >= max_size || (self.pos + buf.len()) > max_size {
 181             io_error::cond.raise(IoError {
 182                 kind: OtherIoError,
 183                 desc: "Trying to write past end of buffer",
 184                 detail: None
 185             });
 186             return;
 187         }
 188
 189         vec::bytes::copy_memory(self.buf.mut_slice_from(self.pos), buf);
 190         self.pos += buf.len();
 191     }
 192 }
 193
 194 // FIXME(#10432)
 195 impl<'a> Seek for BufWriter<'a> {
 196     fn tell(&self) -> u64 { self.pos as u64 }
 197
 198     fn seek(&mut self, pos: i64, style: SeekStyle) {
 199         // compute offset as signed and clamp to prevent overflow
 200         let offset = match style {
 201             SeekSet => { 0 }
 202             SeekEnd => { self.buf.len() }
 203             SeekCur => { self.pos }
 204         } as i64;
 205
 206         self.pos = max(0, offset+pos) as uint;
 207     }
 208 }
 209
 210
 211 /// Reads from a fixed-size byte slice
 212 pub struct BufReader<'a> {
 213     priv buf: &'a [u8],
 214     priv pos: uint
 215 }
 216
 217 impl<'a> BufReader<'a> {
 218     pub fn new<'a>(buf: &'a [u8]) -> BufReader<'a> {
 219         BufReader {
 220             buf: buf,
 221             pos: 0
 222         }
 223     }
 224
 225     /// Tests whether this reader has read all bytes in its buffer.
 226     ///
 227     /// If `true`, then this will no longer return bytes from `read`.
 228     pub fn eof(&self) -> bool { self.pos == self.buf.len() }
 229 }
 230
 231 impl<'a> Reader for BufReader<'a> {
 232     fn read(&mut self, buf: &mut [u8]) -> Option<uint> {
 233         if self.eof() { return None }
 234
 235         let write_len = min(buf.len(), self.buf.len() - self.pos);
 236         {
 237             let input = self.buf.slice(self.pos, self.pos + write_len);
 238             let output = buf.mut_slice(0, write_len);
 239             assert_eq!(input.len(), output.len());
 240             vec::bytes::copy_memory(output, input);
 241         }
 242         self.pos += write_len;
 243         assert!(self.pos <= self.buf.len());
 244
 245         return Some(write_len);
 246      }
 247 }
 248
 249 impl<'a> Seek for BufReader<'a> {
 250     fn tell(&self) -> u64 { self.pos as u64 }
 251
 252     fn seek(&mut self, _pos: i64, _style: SeekStyle) { fail!() }
 253 }
 254
 255 impl<'a> Buffer for BufReader<'a> {
 256     fn fill<'a>(&'a mut self) -> &'a [u8] { self.buf.slice_from(self.pos) }
 257     fn consume(&mut self, amt: uint) { self.pos += amt; }
 258 }
 259
 260 ///Calls a function with a MemWriter and returns
 261 ///the writer's stored vector.
 262 pub fn with_mem_writer(writeFn: |&mut MemWriter|) -> ~[u8] {
 263     let mut writer = MemWriter::new();
 264     writeFn(&mut writer);
 265     writer.unwrap()
 266 }
 267
 268 #[cfg(test)]
 269 mod test {
 270     use prelude::*;
 271     use super::*;
 272     use io::*;
 273
 274     #[test]
 275     fn test_mem_writer() {
 276         let mut writer = MemWriter::new();
 277         assert_eq!(writer.tell(), 0);
 278         writer.write([0]);
 279         assert_eq!(writer.tell(), 1);
 280         writer.write([1, 2, 3]);
 281         writer.write([4, 5, 6, 7]);
 282         assert_eq!(writer.tell(), 8);
 283         assert_eq!(writer.get_ref(), [0, 1, 2, 3, 4, 5, 6, 7]);
 284
 285         writer.seek(0, SeekSet);
 286         assert_eq!(writer.tell(), 0);
 287         writer.write([3, 4]);
 288         assert_eq!(writer.get_ref(), [3, 4, 2, 3, 4, 5, 6, 7]);
 289
 290         writer.seek(1, SeekCur);
 291         writer.write([0, 1]);
 292         assert_eq!(writer.get_ref(), [3, 4, 2, 0, 1, 5, 6, 7]);
 293
 294         writer.seek(-1, SeekEnd);
 295         writer.write([1, 2]);
 296         assert_eq!(writer.get_ref(), [3, 4, 2, 0, 1, 5, 6, 1, 2]);
 297
 298         writer.seek(1, SeekEnd);
 299         writer.write([1]);
 300         assert_eq!(writer.get_ref(), [3, 4, 2, 0, 1, 5, 6, 1, 2, 0, 1]);
 301     }
 302
 303     #[test]
 304     fn test_buf_writer() {
 305         let mut buf = [0 as u8, ..8];
 306         {
 307             let mut writer = BufWriter::new(buf);
 308             assert_eq!(writer.tell(), 0);
 309             writer.write([0]);
 310             assert_eq!(writer.tell(), 1);
 311             writer.write([1, 2, 3]);
 312             writer.write([4, 5, 6, 7]);
 313             assert_eq!(writer.tell(), 8);
 314         }
 315         assert_eq!(buf, [0, 1, 2, 3, 4, 5, 6, 7]);
 316     }
 317
 318     #[test]
 319     fn test_buf_writer_seek() {
 320         let mut buf = [0 as u8, ..8];
 321         {
 322             let mut writer = BufWriter::new(buf);
 323             assert_eq!(writer.tell(), 0);
 324             writer.write([1]);
 325             assert_eq!(writer.tell(), 1);
 326
 327             writer.seek(2, SeekSet);
 328             assert_eq!(writer.tell(), 2);
 329             writer.write([2]);
 330             assert_eq!(writer.tell(), 3);
 331
 332             writer.seek(-2, SeekCur);
 333             assert_eq!(writer.tell(), 1);
 334             writer.write([3]);
 335             assert_eq!(writer.tell(), 2);
 336
 337             writer.seek(-1, SeekEnd);
 338             assert_eq!(writer.tell(), 7);
 339             writer.write([4]);
 340             assert_eq!(writer.tell(), 8);
 341
 342         }
 343         assert_eq!(buf, [1, 3, 2, 0, 0, 0, 0, 4]);
 344     }
 345
 346     #[test]
 347     fn test_buf_writer_error() {
 348         let mut buf = [0 as u8, ..2];
 349         let mut writer = BufWriter::new(buf);
 350         writer.write([0]);
 351
 352         let mut called = false;
 353         io_error::cond.trap(|err| {
 354             assert_eq!(err.kind, OtherIoError);
 355             called = true;
 356         }).inside(|| {
 357             writer.write([0, 0]);
 358         });
 359         assert!(called);
 360     }
 361
 362     #[test]
 363     fn test_mem_reader() {
 364         let mut reader = MemReader::new(~[0, 1, 2, 3, 4, 5, 6, 7]);
 365         let mut buf = [];
 366         assert_eq!(reader.read(buf), Some(0));
 367         assert_eq!(reader.tell(), 0);
 368         let mut buf = [0];
 369         assert_eq!(reader.read(buf), Some(1));
 370         assert_eq!(reader.tell(), 1);
 371         assert_eq!(buf, [0]);
 372         let mut buf = [0, ..4];
 373         assert_eq!(reader.read(buf), Some(4));
 374         assert_eq!(reader.tell(), 5);
 375         assert_eq!(buf, [1, 2, 3, 4]);
 376         assert_eq!(reader.read(buf), Some(3));
 377         assert_eq!(buf.slice(0, 3), [5, 6, 7]);
 378         assert_eq!(reader.read(buf), None);
 379     }
 380
 381     #[test]
 382     fn test_buf_reader() {
 383         let in_buf = ~[0, 1, 2, 3, 4, 5, 6, 7];
 384         let mut reader = BufReader::new(in_buf);
 385         let mut buf = [];
 386         assert_eq!(reader.read(buf), Some(0));
 387         assert_eq!(reader.tell(), 0);
 388         let mut buf = [0];
 389         assert_eq!(reader.read(buf), Some(1));
 390         assert_eq!(reader.tell(), 1);
 391         assert_eq!(buf, [0]);
 392         let mut buf = [0, ..4];
 393         assert_eq!(reader.read(buf), Some(4));
 394         assert_eq!(reader.tell(), 5);
 395         assert_eq!(buf, [1, 2, 3, 4]);
 396         assert_eq!(reader.read(buf), Some(3));
 397         assert_eq!(buf.slice(0, 3), [5, 6, 7]);
 398         assert_eq!(reader.read(buf), None);
 399     }
 400
 401     #[test]
 402     fn test_with_mem_writer() {
 403         let buf = with_mem_writer(|wr| wr.write([1,2,3,4,5,6,7]));
 404         assert_eq!(buf, ~[1,2,3,4,5,6,7]);
 405     }
 406
 407     #[test]
 408     fn test_read_char() {
 409         let b = bytes!("Việt");
 410         let mut r = BufReader::new(b);
 411         assert_eq!(r.read_char(), Some('V'));
 412         assert_eq!(r.read_char(), Some('i'));
 413         assert_eq!(r.read_char(), Some('ệ'));
 414         assert_eq!(r.read_char(), Some('t'));
 415         assert_eq!(r.read_char(), None);
 416     }
 417
 418     #[test]
 419     fn test_read_bad_char() {
 420         let b = bytes!(0x80);
 421         let mut r = BufReader::new(b);
 422         assert_eq!(r.read_char(), None);
 423     }
 424
 425     #[test]
 426     fn test_write_strings() {
 427         let mut writer = MemWriter::new();
 428         writer.write_str("testing");
 429         writer.write_line("testing");
 430         writer.write_str("testing");
 431         let mut r = BufReader::new(writer.get_ref());
 432         assert_eq!(r.read_to_str(), ~"testingtesting\ntesting");
 433     }
 434
 435     #[test]
 436     fn test_write_char() {
 437         let mut writer = MemWriter::new();
 438         writer.write_char('a');
 439         writer.write_char('\n');
 440         writer.write_char('ệ');
 441         let mut r = BufReader::new(writer.get_ref());
 442         assert_eq!(r.read_to_str(), ~"a\nệ");
 443     }
 444
 445     #[test]
 446     fn test_read_whole_string_bad() {
 447         let buf = [0xff];
 448         let mut r = BufReader::new(buf);
 449         match result(|| r.read_to_str()) {
 450             Ok(..) => fail!(),
 451             Err(..) => {}
 452         }
 453     }
 454 }