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