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 use cmp::min;
  14 use container::Container;
  15 use option::None;
  16 use result::{Err, Ok};
  17 use io;
  18 use io::{Reader, Writer, Seek, Buffer, IoError, SeekStyle, IoResult};
  19 use slice;
  20 use slice::{Vector, ImmutableVector, MutableVector};
  21 use vec::Vec;
  22
  23 fn combine(seek: SeekStyle, cur: uint, end: uint, offset: i64) -> IoResult<u64> {
  24     // compute offset as signed and clamp to prevent overflow
  25     let pos = match seek {
  26         SeekSet => 0,
  27         SeekEnd => end,
  28         SeekCur => cur,
  29     } as i64;
  30
  31     if offset + pos < 0 {
  32         Err(IoError {
  33             kind: io::InvalidInput,
  34             desc: "invalid seek to a negative offset",
  35             detail: None
  36         })
  37     } else {
  38         Ok((offset + pos) as u64)
  39     }
  40 }
  41
  42 /// Writes to an owned, growable byte vector
  43 ///
  44 /// # Example
  45 ///
  46 /// ```rust
  47 /// # #![allow(unused_must_use)]
  48 /// use std::io::MemWriter;
  49 ///
  50 /// let mut w = MemWriter::new();
  51 /// w.write([0, 1, 2]);
  52 ///
  53 /// assert_eq!(w.unwrap(), vec!(0, 1, 2));
  54 /// ```
  55 pub struct MemWriter {
  56     buf: Vec<u8>,
  57     pos: uint,
  58 }
  59
  60 impl MemWriter {
  61     /// Create a new `MemWriter`.
  62     pub fn new() -> MemWriter {
  63         MemWriter::with_capacity(128)
  64     }
  65     /// Create a new `MemWriter`, allocating at least `n` bytes for
  66     /// the internal buffer.
  67     pub fn with_capacity(n: uint) -> MemWriter {
  68         MemWriter { buf: Vec::with_capacity(n), pos: 0 }
  69     }
  70
  71     /// Acquires an immutable reference to the underlying buffer of this
  72     /// `MemWriter`.
  73     ///
  74     /// No method is exposed for acquiring a mutable reference to the buffer
  75     /// because it could corrupt the state of this `MemWriter`.
  76     pub fn get_ref<'a>(&'a self) -> &'a [u8] { self.buf.as_slice() }
  77
  78     /// Unwraps this `MemWriter`, returning the underlying buffer
  79     pub fn unwrap(self) -> Vec<u8> { self.buf }
  80 }
  81
  82 impl Writer for MemWriter {
  83     fn write(&mut self, buf: &[u8]) -> IoResult<()> {
  84         // Make sure the internal buffer is as least as big as where we
  85         // currently are
  86         let difference = self.pos as i64 - self.buf.len() as i64;
  87         if difference > 0 {
  88             self.buf.grow(difference as uint, &0);
  89         }
  90
  91         // Figure out what bytes will be used to overwrite what's currently
  92         // there (left), and what will be appended on the end (right)
  93         let cap = self.buf.len() - self.pos;
  94         let (left, right) = if cap <= buf.len() {
  95             (buf.slice_to(cap), buf.slice_from(cap))
  96         } else {
  97             (buf, &[])
  98         };
  99
 100         // Do the necessary writes
 101         if left.len() > 0 {
 102             slice::bytes::copy_memory(self.buf.mut_slice_from(self.pos), left);
 103         }
 104         if right.len() > 0 {
 105             self.buf.push_all(right);
 106         }
 107
 108         // Bump us forward
 109         self.pos += buf.len();
 110         Ok(())
 111     }
 112 }
 113
 114 impl Seek for MemWriter {
 115     fn tell(&self) -> IoResult<u64> { Ok(self.pos as u64) }
 116     fn seek(&mut self, pos: i64, style: SeekStyle) -> IoResult<()> {
 117         let new = try!(combine(style, self.pos, self.buf.len(), pos));
 118         self.pos = new as uint;
 119         Ok(())
 120     }
 121 }
 122
 123 /// Reads from an owned byte vector
 124 ///
 125 /// # Example
 126 ///
 127 /// ```rust
 128 /// # #![allow(unused_must_use)]
 129 /// use std::io::MemReader;
 130 ///
 131 /// let mut r = MemReader::new(vec!(0, 1, 2));
 132 ///
 133 /// assert_eq!(r.read_to_end().unwrap(), vec!(0, 1, 2));
 134 /// ```
 135 pub struct MemReader {
 136     buf: Vec<u8>,
 137     pos: uint
 138 }
 139
 140 impl MemReader {
 141     /// Creates a new `MemReader` which will read the buffer given. The buffer
 142     /// can be re-acquired through `unwrap`
 143     pub fn new(buf: Vec<u8>) -> MemReader {
 144         MemReader {
 145             buf: buf,
 146             pos: 0
 147         }
 148     }
 149
 150     /// Tests whether this reader has read all bytes in its buffer.
 151     ///
 152     /// If `true`, then this will no longer return bytes from `read`.
 153     pub fn eof(&self) -> bool { self.pos >= self.buf.len() }
 154
 155     /// Acquires an immutable reference to the underlying buffer of this
 156     /// `MemReader`.
 157     ///
 158     /// No method is exposed for acquiring a mutable reference to the buffer
 159     /// because it could corrupt the state of this `MemReader`.
 160     pub fn get_ref<'a>(&'a self) -> &'a [u8] { self.buf.as_slice() }
 161
 162     /// Unwraps this `MemReader`, returning the underlying buffer
 163     pub fn unwrap(self) -> Vec<u8> { self.buf }
 164 }
 165
 166 impl Reader for MemReader {
 167     fn read(&mut self, buf: &mut [u8]) -> IoResult<uint> {
 168         if self.eof() { return Err(io::standard_error(io::EndOfFile)) }
 169
 170         let write_len = min(buf.len(), self.buf.len() - self.pos);
 171         {
 172             let input = self.buf.slice(self.pos, self.pos + write_len);
 173             let output = buf.mut_slice(0, write_len);
 174             assert_eq!(input.len(), output.len());
 175             slice::bytes::copy_memory(output, input);
 176         }
 177         self.pos += write_len;
 178         assert!(self.pos <= self.buf.len());
 179
 180         return Ok(write_len);
 181     }
 182 }
 183
 184 impl Seek for MemReader {
 185     fn tell(&self) -> IoResult<u64> { Ok(self.pos as u64) }
 186     fn seek(&mut self, pos: i64, style: SeekStyle) -> IoResult<()> {
 187         let new = try!(combine(style, self.pos, self.buf.len(), pos));
 188         self.pos = new as uint;
 189         Ok(())
 190     }
 191 }
 192
 193 impl Buffer for MemReader {
 194     fn fill_buf<'a>(&'a mut self) -> IoResult<&'a [u8]> {
 195         if self.pos < self.buf.len() {
 196             Ok(self.buf.slice_from(self.pos))
 197         } else {
 198             Err(io::standard_error(io::EndOfFile))
 199         }
 200     }
 201     fn consume(&mut self, amt: uint) { self.pos += amt; }
 202 }
 203
 204 /// Writes to a fixed-size byte slice
 205 ///
 206 /// If a write will not fit in the buffer, it returns an error and does not
 207 /// write any data.
 208 ///
 209 /// # Example
 210 ///
 211 /// ```rust
 212 /// # #![allow(unused_must_use)]
 213 /// use std::io::BufWriter;
 214 ///
 215 /// let mut buf = [0, ..4];
 216 /// {
 217 ///     let mut w = BufWriter::new(buf);
 218 ///     w.write([0, 1, 2]);
 219 /// }
 220 /// assert!(buf == [0, 1, 2, 0]);
 221 /// ```
 222 pub struct BufWriter<'a> {
 223     buf: &'a mut [u8],
 224     pos: uint
 225 }
 226
 227 impl<'a> BufWriter<'a> {
 228     /// Creates a new `BufWriter` which will wrap the specified buffer. The
 229     /// writer initially starts at position 0.
 230     pub fn new<'a>(buf: &'a mut [u8]) -> BufWriter<'a> {
 231         BufWriter {
 232             buf: buf,
 233             pos: 0
 234         }
 235     }
 236 }
 237
 238 impl<'a> Writer for BufWriter<'a> {
 239     fn write(&mut self, buf: &[u8]) -> IoResult<()> {
 240         // return an error if the entire write does not fit in the buffer
 241         let max_size = self.buf.len();
 242         if self.pos >= max_size || (self.pos + buf.len()) > max_size {
 243             return Err(IoError {
 244                 kind: io::OtherIoError,
 245                 desc: "Trying to write past end of buffer",
 246                 detail: None
 247             })
 248         }
 249
 250         slice::bytes::copy_memory(self.buf.mut_slice_from(self.pos), buf);
 251         self.pos += buf.len();
 252         Ok(())
 253     }
 254 }
 255
 256 impl<'a> Seek for BufWriter<'a> {
 257     fn tell(&self) -> IoResult<u64> { Ok(self.pos as u64) }
 258     fn seek(&mut self, pos: i64, style: SeekStyle) -> IoResult<()> {
 259         let new = try!(combine(style, self.pos, self.buf.len(), pos));
 260         self.pos = new as uint;
 261         Ok(())
 262     }
 263 }
 264
 265 /// Reads from a fixed-size byte slice
 266 ///
 267 /// # Example
 268 ///
 269 /// ```rust
 270 /// # #![allow(unused_must_use)]
 271 /// use std::io::BufReader;
 272 ///
 273 /// let mut buf = [0, 1, 2, 3];
 274 /// let mut r = BufReader::new(buf);
 275 ///
 276 /// assert_eq!(r.read_to_end().unwrap(), vec!(0, 1, 2, 3));
 277 /// ```
 278 pub struct BufReader<'a> {
 279     buf: &'a [u8],
 280     pos: uint
 281 }
 282
 283 impl<'a> BufReader<'a> {
 284     /// Creates a new buffered reader which will read the specified buffer
 285     pub fn new<'a>(buf: &'a [u8]) -> BufReader<'a> {
 286         BufReader {
 287             buf: buf,
 288             pos: 0
 289         }
 290     }
 291
 292     /// Tests whether this reader has read all bytes in its buffer.
 293     ///
 294     /// If `true`, then this will no longer return bytes from `read`.
 295     pub fn eof(&self) -> bool { self.pos >= self.buf.len() }
 296 }
 297
 298 impl<'a> Reader for BufReader<'a> {
 299     fn read(&mut self, buf: &mut [u8]) -> IoResult<uint> {
 300         if self.eof() { return Err(io::standard_error(io::EndOfFile)) }
 301
 302         let write_len = min(buf.len(), self.buf.len() - self.pos);
 303         {
 304             let input = self.buf.slice(self.pos, self.pos + write_len);
 305             let output = buf.mut_slice(0, write_len);
 306             assert_eq!(input.len(), output.len());
 307             slice::bytes::copy_memory(output, input);
 308         }
 309         self.pos += write_len;
 310         assert!(self.pos <= self.buf.len());
 311
 312         return Ok(write_len);
 313      }
 314 }
 315
 316 impl<'a> Seek for BufReader<'a> {
 317     fn tell(&self) -> IoResult<u64> { Ok(self.pos as u64) }
 318     fn seek(&mut self, pos: i64, style: SeekStyle) -> IoResult<()> {
 319         let new = try!(combine(style, self.pos, self.buf.len(), pos));
 320         self.pos = new as uint;
 321         Ok(())
 322     }
 323 }
 324
 325 impl<'a> Buffer for BufReader<'a> {
 326     fn fill_buf<'a>(&'a mut self) -> IoResult<&'a [u8]> {
 327         if self.pos < self.buf.len() {
 328             Ok(self.buf.slice_from(self.pos))
 329         } else {
 330             Err(io::standard_error(io::EndOfFile))
 331         }
 332     }
 333     fn consume(&mut self, amt: uint) { self.pos += amt; }
 334 }
 335
 336 #[cfg(test)]
 337 mod test {
 338     use prelude::*;
 339     use super::*;
 340     use io::*;
 341     use io;
 342     use str::StrSlice;
 343
 344     #[test]
 345     fn test_mem_writer() {
 346         let mut writer = MemWriter::new();
 347         assert_eq!(writer.tell(), Ok(0));
 348         writer.write([0]).unwrap();
 349         assert_eq!(writer.tell(), Ok(1));
 350         writer.write([1, 2, 3]).unwrap();
 351         writer.write([4, 5, 6, 7]).unwrap();
 352         assert_eq!(writer.tell(), Ok(8));
 353         assert_eq!(writer.get_ref(), &[0, 1, 2, 3, 4, 5, 6, 7]);
 354
 355         writer.seek(0, SeekSet).unwrap();
 356         assert_eq!(writer.tell(), Ok(0));
 357         writer.write([3, 4]).unwrap();
 358         assert_eq!(writer.get_ref(), &[3, 4, 2, 3, 4, 5, 6, 7]);
 359
 360         writer.seek(1, SeekCur).unwrap();
 361         writer.write([0, 1]).unwrap();
 362         assert_eq!(writer.get_ref(), &[3, 4, 2, 0, 1, 5, 6, 7]);
 363
 364         writer.seek(-1, SeekEnd).unwrap();
 365         writer.write([1, 2]).unwrap();
 366         assert_eq!(writer.get_ref(), &[3, 4, 2, 0, 1, 5, 6, 1, 2]);
 367
 368         writer.seek(1, SeekEnd).unwrap();
 369         writer.write([1]).unwrap();
 370         assert_eq!(writer.get_ref(), &[3, 4, 2, 0, 1, 5, 6, 1, 2, 0, 1]);
 371     }
 372
 373     #[test]
 374     fn test_buf_writer() {
 375         let mut buf = [0 as u8, ..8];
 376         {
 377             let mut writer = BufWriter::new(buf);
 378             assert_eq!(writer.tell(), Ok(0));
 379             writer.write([0]).unwrap();
 380             assert_eq!(writer.tell(), Ok(1));
 381             writer.write([1, 2, 3]).unwrap();
 382             writer.write([4, 5, 6, 7]).unwrap();
 383             assert_eq!(writer.tell(), Ok(8));
 384         }
 385         assert_eq!(buf.as_slice(), &[0, 1, 2, 3, 4, 5, 6, 7]);
 386     }
 387
 388     #[test]
 389     fn test_buf_writer_seek() {
 390         let mut buf = [0 as u8, ..8];
 391         {
 392             let mut writer = BufWriter::new(buf);
 393             assert_eq!(writer.tell(), Ok(0));
 394             writer.write([1]).unwrap();
 395             assert_eq!(writer.tell(), Ok(1));
 396
 397             writer.seek(2, SeekSet).unwrap();
 398             assert_eq!(writer.tell(), Ok(2));
 399             writer.write([2]).unwrap();
 400             assert_eq!(writer.tell(), Ok(3));
 401
 402             writer.seek(-2, SeekCur).unwrap();
 403             assert_eq!(writer.tell(), Ok(1));
 404             writer.write([3]).unwrap();
 405             assert_eq!(writer.tell(), Ok(2));
 406
 407             writer.seek(-1, SeekEnd).unwrap();
 408             assert_eq!(writer.tell(), Ok(7));
 409             writer.write([4]).unwrap();
 410             assert_eq!(writer.tell(), Ok(8));
 411
 412         }
 413         assert_eq!(buf.as_slice(), &[1, 3, 2, 0, 0, 0, 0, 4]);
 414     }
 415
 416     #[test]
 417     fn test_buf_writer_error() {
 418         let mut buf = [0 as u8, ..2];
 419         let mut writer = BufWriter::new(buf);
 420         writer.write([0]).unwrap();
 421
 422         match writer.write([0, 0]) {
 423             Ok(..) => fail!(),
 424             Err(e) => assert_eq!(e.kind, io::OtherIoError),
 425         }
 426     }
 427
 428     #[test]
 429     fn test_mem_reader() {
 430         let mut reader = MemReader::new(vec!(0, 1, 2, 3, 4, 5, 6, 7));
 431         let mut buf = [];
 432         assert_eq!(reader.read(buf), Ok(0));
 433         assert_eq!(reader.tell(), Ok(0));
 434         let mut buf = [0];
 435         assert_eq!(reader.read(buf), Ok(1));
 436         assert_eq!(reader.tell(), Ok(1));
 437         assert_eq!(buf.as_slice(), &[0]);
 438         let mut buf = [0, ..4];
 439         assert_eq!(reader.read(buf), Ok(4));
 440         assert_eq!(reader.tell(), Ok(5));
 441         assert_eq!(buf.as_slice(), &[1, 2, 3, 4]);
 442         assert_eq!(reader.read(buf), Ok(3));
 443         assert_eq!(buf.slice(0, 3), &[5, 6, 7]);
 444         assert!(reader.read(buf).is_err());
 445         let mut reader = MemReader::new(vec!(0, 1, 2, 3, 4, 5, 6, 7));
 446         assert_eq!(reader.read_until(3).unwrap(), vec!(0, 1, 2, 3));
 447         assert_eq!(reader.read_until(3).unwrap(), vec!(4, 5, 6, 7));
 448         assert!(reader.read(buf).is_err());
 449     }
 450
 451     #[test]
 452     fn test_buf_reader() {
 453         let in_buf = ~[0, 1, 2, 3, 4, 5, 6, 7];
 454         let mut reader = BufReader::new(in_buf);
 455         let mut buf = [];
 456         assert_eq!(reader.read(buf), Ok(0));
 457         assert_eq!(reader.tell(), Ok(0));
 458         let mut buf = [0];
 459         assert_eq!(reader.read(buf), Ok(1));
 460         assert_eq!(reader.tell(), Ok(1));
 461         assert_eq!(buf.as_slice(), &[0]);
 462         let mut buf = [0, ..4];
 463         assert_eq!(reader.read(buf), Ok(4));
 464         assert_eq!(reader.tell(), Ok(5));
 465         assert_eq!(buf.as_slice(), &[1, 2, 3, 4]);
 466         assert_eq!(reader.read(buf), Ok(3));
 467         assert_eq!(buf.slice(0, 3), &[5, 6, 7]);
 468         assert!(reader.read(buf).is_err());
 469         let mut reader = BufReader::new(in_buf);
 470         assert_eq!(reader.read_until(3).unwrap(), vec!(0, 1, 2, 3));
 471         assert_eq!(reader.read_until(3).unwrap(), vec!(4, 5, 6, 7));
 472         assert!(reader.read(buf).is_err());
 473     }
 474
 475     #[test]
 476     fn test_read_char() {
 477         let b = bytes!("Việt");
 478         let mut r = BufReader::new(b);
 479         assert_eq!(r.read_char(), Ok('V'));
 480         assert_eq!(r.read_char(), Ok('i'));
 481         assert_eq!(r.read_char(), Ok('ệ'));
 482         assert_eq!(r.read_char(), Ok('t'));
 483         assert!(r.read_char().is_err());
 484     }
 485
 486     #[test]
 487     fn test_read_bad_char() {
 488         let b = bytes!(0x80);
 489         let mut r = BufReader::new(b);
 490         assert!(r.read_char().is_err());
 491     }
 492
 493     #[test]
 494     fn test_write_strings() {
 495         let mut writer = MemWriter::new();
 496         writer.write_str("testing").unwrap();
 497         writer.write_line("testing").unwrap();
 498         writer.write_str("testing").unwrap();
 499         let mut r = BufReader::new(writer.get_ref());
 500         assert_eq!(r.read_to_str().unwrap(), "testingtesting\ntesting".to_owned());
 501     }
 502
 503     #[test]
 504     fn test_write_char() {
 505         let mut writer = MemWriter::new();
 506         writer.write_char('a').unwrap();
 507         writer.write_char('\n').unwrap();
 508         writer.write_char('ệ').unwrap();
 509         let mut r = BufReader::new(writer.get_ref());
 510         assert_eq!(r.read_to_str().unwrap(), "a\nệ".to_owned());
 511     }
 512
 513     #[test]
 514     fn test_read_whole_string_bad() {
 515         let buf = [0xff];
 516         let mut r = BufReader::new(buf);
 517         match r.read_to_str() {
 518             Ok(..) => fail!(),
 519             Err(..) => {}
 520         }
 521     }
 522
 523     #[test]
 524     fn seek_past_end() {
 525         let buf = [0xff];
 526         let mut r = BufReader::new(buf);
 527         r.seek(10, SeekSet).unwrap();
 528         assert!(r.read(&mut []).is_err());
 529
 530         let mut r = MemReader::new(vec!(10));
 531         r.seek(10, SeekSet).unwrap();
 532         assert!(r.read(&mut []).is_err());
 533
 534         let mut r = MemWriter::new();
 535         r.seek(10, SeekSet).unwrap();
 536         assert!(r.write([3]).is_ok());
 537
 538         let mut buf = [0];
 539         let mut r = BufWriter::new(buf);
 540         r.seek(10, SeekSet).unwrap();
 541         assert!(r.write([3]).is_err());
 542     }
 543
 544     #[test]
 545     fn seek_before_0() {
 546         let buf = [0xff];
 547         let mut r = BufReader::new(buf);
 548         assert!(r.seek(-1, SeekSet).is_err());
 549
 550         let mut r = MemReader::new(vec!(10));
 551         assert!(r.seek(-1, SeekSet).is_err());
 552
 553         let mut r = MemWriter::new();
 554         assert!(r.seek(-1, SeekSet).is_err());
 555
 556         let mut buf = [0];
 557         let mut r = BufWriter::new(buf);
 558         assert!(r.seek(-1, SeekSet).is_err());
 559     }
 560
 561     #[test]
 562     fn io_fill() {
 563         let mut r = MemReader::new(vec!(1, 2, 3, 4, 5, 6, 7, 8));
 564         let mut buf = [0, ..3];
 565         assert_eq!(r.fill(buf), Ok(()));
 566         assert_eq!(buf.as_slice(), &[1, 2, 3]);
 567         assert_eq!(r.fill(buf.mut_slice_to(0)), Ok(()));
 568         assert_eq!(buf.as_slice(), &[1, 2, 3]);
 569         assert_eq!(r.fill(buf), Ok(()));
 570         assert_eq!(buf.as_slice(), &[4, 5, 6]);
 571         assert!(r.fill(buf).is_err());
 572         assert_eq!(buf.as_slice(), &[7, 8, 6]);
 573     }
 574 }