src/libstd/sys/common/io.rs

   1 // Copyright 2015 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::v1::*;
  11 use io;
  12 use io::ErrorKind;
  13 use io::Read;
  14 use slice::from_raw_parts_mut;
  15
  16 // Provides read_to_end functionality over an uninitialized buffer.
  17 // This function is unsafe because it calls the underlying
  18 // read function with a slice into uninitialized memory. The default
  19 // implementation of read_to_end for readers will zero out new memory in
  20 // the buf before passing it to read, but avoiding this zero can often
  21 // lead to a fairly significant performance win.
  22 //
  23 // Implementations using this method have to adhere to two guarantees:
  24 //  *  The implementation of read never reads the buffer provided.
  25 //  *  The implementation of read correctly reports how many bytes were written.
  26 pub unsafe fn read_to_end_uninitialized(r: &mut Read, buf: &mut Vec<u8>) -> io::Result<usize> {
  27
  28     let start_len = buf.len();
  29     buf.reserve(16);
  30
  31     // Always try to read into the empty space of the vector (from the length to the capacity).
  32     // If the vector ever fills up then we reserve an extra byte which should trigger the normal
  33     // reallocation routines for the vector, which will likely double the size.
  34     //
  35     // This function is similar to the read_to_end function in std::io, but the logic about
  36     // reservations and slicing is different enough that this is duplicated here.
  37     loop {
  38         if buf.len() == buf.capacity() {
  39             buf.reserve(1);
  40         }
  41
  42         let buf_slice = from_raw_parts_mut(buf.as_mut_ptr().offset(buf.len() as isize),
  43                                            buf.capacity() - buf.len());
  44
  45         match r.read(buf_slice) {
  46             Ok(0) => { return Ok(buf.len() - start_len); }
  47             Ok(n) => { let len = buf.len() + n; buf.set_len(len); },
  48             Err(ref e) if e.kind() == ErrorKind::Interrupted => { }
  49             Err(e) => { return Err(e); }
  50         }
  51     }
  52 }
  53
  54 #[cfg(test)]
  55 mod tests {
  56     use prelude::v1::*;
  57     use io::prelude::*;
  58     use super::*;
  59     use io;
  60     use io::{ErrorKind, Take, Repeat, repeat};
  61     use test;
  62     use slice::from_raw_parts;
  63
  64     struct ErrorRepeat {
  65         lr: Take<Repeat>
  66     }
  67
  68     fn error_repeat(byte: u8, limit: u64) -> ErrorRepeat {
  69         ErrorRepeat { lr: repeat(byte).take(limit) }
  70     }
  71
  72     impl Read for ErrorRepeat {
  73         fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
  74             let ret = self.lr.read(buf);
  75             if let Ok(0) = ret {
  76                 return Err(io::Error::new(ErrorKind::Other, ""))
  77             }
  78             ret
  79         }
  80     }
  81
  82     fn init_vec_data() -> Vec<u8> {
  83         let mut vec = vec![10u8; 200];
  84         unsafe { vec.set_len(0); }
  85         vec
  86     }
  87
  88     fn assert_all_eq(buf: &[u8], value: u8) {
  89         for n in buf {
  90             assert_eq!(*n, value);
  91         }
  92     }
  93
  94     fn validate(buf: &Vec<u8>, good_read_len: usize) {
  95         assert_all_eq(buf, 1u8);
  96         let cap = buf.capacity();
  97         let end_slice = unsafe { from_raw_parts(buf.as_ptr().offset(good_read_len as isize),
  98                                                     cap - good_read_len) };
  99         assert_all_eq(end_slice, 10u8);
 100     }
 101
 102     #[test]
 103     fn read_to_end_uninit_error() {
 104         let mut er = error_repeat(1,100);
 105         let mut vec = init_vec_data();
 106         if let Err(_) = unsafe { read_to_end_uninitialized(&mut er, &mut vec) } {
 107             validate(&vec, 100);
 108         } else {
 109             assert!(false);
 110         }
 111     }
 112
 113     #[test]
 114     fn read_to_end_uninit_zero_len_vec() {
 115         let mut er = repeat(1).take(100);
 116         let mut vec = Vec::new();
 117         let n = unsafe{ read_to_end_uninitialized(&mut er, &mut vec).unwrap() };
 118         assert_all_eq(&vec, 1u8);
 119         assert_eq!(vec.len(), n);
 120     }
 121
 122     #[test]
 123     fn read_to_end_uninit_good() {
 124         let mut er = repeat(1).take(100);
 125         let mut vec = init_vec_data();
 126         let n = unsafe{ read_to_end_uninitialized(&mut er, &mut vec).unwrap() };
 127         validate(&vec, 100);
 128         assert_eq!(vec.len(), n);
 129     }
 130
 131     #[bench]
 132     fn bench_uninitialized(b: &mut test::Bencher) {
 133         b.iter(|| {
 134             let mut lr = repeat(1).take(10000000);
 135             let mut vec = Vec::with_capacity(1024);
 136             unsafe { read_to_end_uninitialized(&mut lr, &mut vec) }
 137         });
 138     }
 139 }