src/librbml/io.rs

   1 // Copyright 2012-2014 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 use std::io::{IoError, IoResult, SeekStyle};
  12 use std::io;
  13 use std::slice;
  14
  15 static BUF_CAPACITY: uint = 128;
  16
  17 fn combine(seek: SeekStyle, cur: uint, end: uint, offset: i64) -> IoResult<u64> {
  18     // compute offset as signed and clamp to prevent overflow
  19     let pos = match seek {
  20         io::SeekSet => 0,
  21         io::SeekEnd => end,
  22         io::SeekCur => cur,
  23     } as i64;
  24
  25     if offset + pos < 0 {
  26         Err(IoError {
  27             kind: io::InvalidInput,
  28             desc: "invalid seek to a negative offset",
  29             detail: None
  30         })
  31     } else {
  32         Ok((offset + pos) as u64)
  33     }
  34 }
  35
  36 /// Writes to an owned, growable byte vector that supports seeking.
  37 ///
  38 /// # Example
  39 ///
  40 /// ```rust
  41 /// # #![allow(unused_must_use)]
  42 /// use rbml::io::SeekableMemWriter;
  43 ///
  44 /// let mut w = SeekableMemWriter::new();
  45 /// w.write([0, 1, 2]);
  46 ///
  47 /// assert_eq!(w.unwrap(), vec!(0, 1, 2));
  48 /// ```
  49 pub struct SeekableMemWriter {
  50     buf: Vec<u8>,
  51     pos: uint,
  52 }
  53
  54 impl SeekableMemWriter {
  55     /// Create a new `SeekableMemWriter`.
  56     #[inline]
  57     pub fn new() -> SeekableMemWriter {
  58         SeekableMemWriter::with_capacity(BUF_CAPACITY)
  59     }
  60     /// Create a new `SeekableMemWriter`, allocating at least `n` bytes for
  61     /// the internal buffer.
  62     #[inline]
  63     pub fn with_capacity(n: uint) -> SeekableMemWriter {
  64         SeekableMemWriter { buf: Vec::with_capacity(n), pos: 0 }
  65     }
  66
  67     /// Acquires an immutable reference to the underlying buffer of this
  68     /// `SeekableMemWriter`.
  69     ///
  70     /// No method is exposed for acquiring a mutable reference to the buffer
  71     /// because it could corrupt the state of this `MemWriter`.
  72     #[inline]
  73     pub fn get_ref<'a>(&'a self) -> &'a [u8] { self.buf.as_slice() }
  74
  75     /// Unwraps this `SeekableMemWriter`, returning the underlying buffer
  76     #[inline]
  77     pub fn unwrap(self) -> Vec<u8> { self.buf }
  78 }
  79
  80 impl Writer for SeekableMemWriter {
  81     #[inline]
  82     fn write(&mut self, buf: &[u8]) -> IoResult<()> {
  83         if self.pos == self.buf.len() {
  84             self.buf.push_all(buf)
  85         } else {
  86             // Make sure the internal buffer is as least as big as where we
  87             // currently are
  88             let difference = self.pos as i64 - self.buf.len() as i64;
  89             if difference > 0 {
  90                 self.buf.grow(difference as uint, &0);
  91             }
  92
  93             // Figure out what bytes will be used to overwrite what's currently
  94             // there (left), and what will be appended on the end (right)
  95             let cap = self.buf.len() - self.pos;
  96             let (left, right) = if cap <= buf.len() {
  97                 (buf.slice_to(cap), buf.slice_from(cap))
  98             } else {
  99                 let result: (_, &[_]) = (buf, &[]);
 100                 result
 101             };
 102
 103             // Do the necessary writes
 104             if left.len() > 0 {
 105                 slice::bytes::copy_memory(self.buf.mut_slice_from(self.pos), left);
 106             }
 107             if right.len() > 0 {
 108                 self.buf.push_all(right);
 109             }
 110         }
 111
 112         // Bump us forward
 113         self.pos += buf.len();
 114         Ok(())
 115     }
 116 }
 117
 118 impl Seek for SeekableMemWriter {
 119     #[inline]
 120     fn tell(&self) -> IoResult<u64> { Ok(self.pos as u64) }
 121
 122     #[inline]
 123     fn seek(&mut self, pos: i64, style: SeekStyle) -> IoResult<()> {
 124         let new = try!(combine(style, self.pos, self.buf.len(), pos));
 125         self.pos = new as uint;
 126         Ok(())
 127     }
 128 }
 129
 130 #[cfg(test)]
 131 mod tests {
 132     extern crate test;
 133     use super::SeekableMemWriter;
 134     use std::io;
 135     use test::Bencher;
 136
 137     #[test]
 138     fn test_seekable_mem_writer() {
 139         let mut writer = SeekableMemWriter::new();
 140         assert_eq!(writer.tell(), Ok(0));
 141         writer.write([0]).unwrap();
 142         assert_eq!(writer.tell(), Ok(1));
 143         writer.write([1, 2, 3]).unwrap();
 144         writer.write([4, 5, 6, 7]).unwrap();
 145         assert_eq!(writer.tell(), Ok(8));
 146         let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7];
 147         assert_eq!(writer.get_ref(), b);
 148
 149         writer.seek(0, io::SeekSet).unwrap();
 150         assert_eq!(writer.tell(), Ok(0));
 151         writer.write([3, 4]).unwrap();
 152         let b: &[_] = &[3, 4, 2, 3, 4, 5, 6, 7];
 153         assert_eq!(writer.get_ref(), b);
 154
 155         writer.seek(1, io::SeekCur).unwrap();
 156         writer.write([0, 1]).unwrap();
 157         let b: &[_] = &[3, 4, 2, 0, 1, 5, 6, 7];
 158         assert_eq!(writer.get_ref(), b);
 159
 160         writer.seek(-1, io::SeekEnd).unwrap();
 161         writer.write([1, 2]).unwrap();
 162         let b: &[_] = &[3, 4, 2, 0, 1, 5, 6, 1, 2];
 163         assert_eq!(writer.get_ref(), b);
 164
 165         writer.seek(1, io::SeekEnd).unwrap();
 166         writer.write([1]).unwrap();
 167         let b: &[_] = &[3, 4, 2, 0, 1, 5, 6, 1, 2, 0, 1];
 168         assert_eq!(writer.get_ref(), b);
 169     }
 170
 171     #[test]
 172     fn seek_past_end() {
 173         let mut r = SeekableMemWriter::new();
 174         r.seek(10, io::SeekSet).unwrap();
 175         assert!(r.write([3]).is_ok());
 176     }
 177
 178     #[test]
 179     fn seek_before_0() {
 180         let mut r = SeekableMemWriter::new();
 181         assert!(r.seek(-1, io::SeekSet).is_err());
 182     }
 183
 184     fn do_bench_seekable_mem_writer(b: &mut Bencher, times: uint, len: uint) {
 185         let src: Vec<u8> = Vec::from_elem(len, 5);
 186
 187         b.bytes = (times * len) as u64;
 188         b.iter(|| {
 189             let mut wr = SeekableMemWriter::new();
 190             for _ in range(0, times) {
 191                 wr.write(src.as_slice()).unwrap();
 192             }
 193
 194             let v = wr.unwrap();
 195             assert_eq!(v.len(), times * len);
 196             assert!(v.iter().all(|x| *x == 5));
 197         });
 198     }
 199
 200     #[bench]
 201     fn bench_seekable_mem_writer_001_0000(b: &mut Bencher) {
 202         do_bench_seekable_mem_writer(b, 1, 0)
 203     }
 204
 205     #[bench]
 206     fn bench_seekable_mem_writer_001_0010(b: &mut Bencher) {
 207         do_bench_seekable_mem_writer(b, 1, 10)
 208     }
 209
 210     #[bench]
 211     fn bench_seekable_mem_writer_001_0100(b: &mut Bencher) {
 212         do_bench_seekable_mem_writer(b, 1, 100)
 213     }
 214
 215     #[bench]
 216     fn bench_seekable_mem_writer_001_1000(b: &mut Bencher) {
 217         do_bench_seekable_mem_writer(b, 1, 1000)
 218     }
 219
 220     #[bench]
 221     fn bench_seekable_mem_writer_100_0000(b: &mut Bencher) {
 222         do_bench_seekable_mem_writer(b, 100, 0)
 223     }
 224
 225     #[bench]
 226     fn bench_seekable_mem_writer_100_0010(b: &mut Bencher) {
 227         do_bench_seekable_mem_writer(b, 100, 10)
 228     }
 229
 230     #[bench]
 231     fn bench_seekable_mem_writer_100_0100(b: &mut Bencher) {
 232         do_bench_seekable_mem_writer(b, 100, 100)
 233     }
 234
 235     #[bench]
 236     fn bench_seekable_mem_writer_100_1000(b: &mut Bencher) {
 237         do_bench_seekable_mem_writer(b, 100, 1000)
 238     }
 239 }