Merge pull request #410 from hyperium/buf-speed
perf(buffer): pull in std::io::BufReader improvements
This commit is contained in:
		
							
								
								
									
										132
									
								
								src/buffer.rs
									
									
									
									
									
								
							
							
						
						
									
										132
									
								
								src/buffer.rs
									
									
									
									
									
								
							| @@ -1,115 +1,143 @@ | |||||||
| use std::cmp; | use std::cmp; | ||||||
| use std::iter; | use std::iter; | ||||||
| use std::io::{self, Read, BufRead, Cursor}; | use std::io::{self, Read, BufRead}; | ||||||
|  |  | ||||||
| pub struct BufReader<R> { | pub struct BufReader<R> { | ||||||
|     buf: Cursor<Vec<u8>>, |     inner: R, | ||||||
|     inner: R |     buf: Vec<u8>, | ||||||
|  |     pos: usize, | ||||||
|  |     cap: usize, | ||||||
| } | } | ||||||
|  |  | ||||||
| const INIT_BUFFER_SIZE: usize = 4096; | const INIT_BUFFER_SIZE: usize = 4096; | ||||||
| const MAX_BUFFER_SIZE: usize = 8192 + 4096 * 100; | const MAX_BUFFER_SIZE: usize = 8192 + 4096 * 100; | ||||||
|  |  | ||||||
| impl<R: Read> BufReader<R> { | impl<R: Read> BufReader<R> { | ||||||
|  |     #[inline] | ||||||
|     pub fn new(rdr: R) -> BufReader<R> { |     pub fn new(rdr: R) -> BufReader<R> { | ||||||
|         BufReader::with_capacity(rdr, INIT_BUFFER_SIZE) |         BufReader::with_capacity(rdr, INIT_BUFFER_SIZE) | ||||||
|     } |     } | ||||||
|  |  | ||||||
|  |     #[inline] | ||||||
|     pub fn with_capacity(rdr: R, cap: usize) -> BufReader<R> { |     pub fn with_capacity(rdr: R, cap: usize) -> BufReader<R> { | ||||||
|  |         let mut buf = Vec::with_capacity(cap); | ||||||
|  |         buf.extend(iter::repeat(0).take(cap)); | ||||||
|         BufReader { |         BufReader { | ||||||
|             buf: Cursor::new(Vec::with_capacity(cap)), |             inner: rdr, | ||||||
|             inner: rdr |             buf: buf, | ||||||
|  |             pos: 0, | ||||||
|  |             cap: 0, | ||||||
|         } |         } | ||||||
|     } |     } | ||||||
|  |  | ||||||
|  |     #[inline] | ||||||
|     pub fn get_ref(&self) -> &R { &self.inner } |     pub fn get_ref(&self) -> &R { &self.inner } | ||||||
|  |  | ||||||
|  |     #[inline] | ||||||
|     pub fn get_mut(&mut self) -> &mut R { &mut self.inner } |     pub fn get_mut(&mut self) -> &mut R { &mut self.inner } | ||||||
|  |  | ||||||
|  |     #[inline] | ||||||
|     pub fn get_buf(&self) -> &[u8] { |     pub fn get_buf(&self) -> &[u8] { | ||||||
|         let pos = self.buf.position() as usize; |         if self.pos < self.cap { | ||||||
|         if pos < self.buf.get_ref().len() { |             debug!("slicing {:?}", (self.pos, self.cap, self.buf.len())); | ||||||
|             &self.buf.get_ref()[pos..] |             &self.buf[self.pos..self.cap] | ||||||
|         } else { |         } else { | ||||||
|             &[] |             &[] | ||||||
|         } |         } | ||||||
|     } |     } | ||||||
|  |  | ||||||
|  |     #[inline] | ||||||
|     pub fn into_inner(self) -> R { self.inner } |     pub fn into_inner(self) -> R { self.inner } | ||||||
|  |  | ||||||
|  |     #[inline] | ||||||
|     pub fn read_into_buf(&mut self) -> io::Result<usize> { |     pub fn read_into_buf(&mut self) -> io::Result<usize> { | ||||||
|         let v = self.buf.get_mut(); |         self.maybe_reserve(); | ||||||
|         reserve(v); |         let v = &mut self.buf; | ||||||
|         let inner = &mut self.inner; |         if self.cap < v.capacity() { | ||||||
|         with_end_to_cap(v, |b| inner.read(b)) |             let nread = try!(self.inner.read(&mut v[self.cap..])); | ||||||
|  |             self.cap += nread; | ||||||
|  |             Ok(nread) | ||||||
|  |         } else { | ||||||
|  |             Ok(0) | ||||||
|  |         } | ||||||
|  |     } | ||||||
|  |  | ||||||
|  |     #[inline] | ||||||
|  |     fn maybe_reserve(&mut self) { | ||||||
|  |         let cap = self.buf.capacity(); | ||||||
|  |         if self.cap == cap { | ||||||
|  |             self.buf.reserve(cmp::min(cap * 4, MAX_BUFFER_SIZE) - cap); | ||||||
|  |         } | ||||||
|     } |     } | ||||||
| } | } | ||||||
|  |  | ||||||
| impl<R: Read> Read for BufReader<R> { | impl<R: Read> Read for BufReader<R> { | ||||||
|     fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> { |     fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> { | ||||||
|         if self.buf.get_ref().len() == self.buf.position() as usize && |         if self.cap == self.pos && buf.len() >= self.buf.len() { | ||||||
|             buf.len() >= self.buf.get_ref().capacity() { |  | ||||||
|             return self.inner.read(buf); |             return self.inner.read(buf); | ||||||
|         } |         } | ||||||
|         try!(self.fill_buf()); |         let nread = { | ||||||
|         self.buf.read(buf) |            let mut rem = try!(self.fill_buf()); | ||||||
|  |            try!(rem.read(buf)) | ||||||
|  |         }; | ||||||
|  |         self.consume(nread); | ||||||
|  |         Ok(nread) | ||||||
|     } |     } | ||||||
| } | } | ||||||
|  |  | ||||||
| impl<R: Read> BufRead for BufReader<R> { | impl<R: Read> BufRead for BufReader<R> { | ||||||
|     fn fill_buf(&mut self) -> io::Result<&[u8]> { |     fn fill_buf(&mut self) -> io::Result<&[u8]> { | ||||||
|          if self.buf.position() as usize == self.buf.get_ref().len() { |         if self.pos == self.cap { | ||||||
|             self.buf.set_position(0); |             self.cap = try!(self.inner.read(&mut self.buf)); | ||||||
|             let v = self.buf.get_mut(); |             self.pos = 0; | ||||||
|             v.truncate(0); |  | ||||||
|             let inner = &mut self.inner; |  | ||||||
|             try!(with_end_to_cap(v, |b| inner.read(b))); |  | ||||||
|         } |         } | ||||||
|          self.buf.fill_buf() |         Ok(&self.buf[self.pos..self.cap]) | ||||||
|     } |     } | ||||||
|  |  | ||||||
|  |     #[inline] | ||||||
|     fn consume(&mut self, amt: usize) { |     fn consume(&mut self, amt: usize) { | ||||||
|         self.buf.consume(amt) |         self.pos = cmp::min(self.pos + amt, self.cap); | ||||||
|  |         if self.pos == self.cap { | ||||||
|  |             self.pos = 0; | ||||||
|  |             self.cap = 0; | ||||||
|         } |         } | ||||||
| } |  | ||||||
|  |  | ||||||
| fn with_end_to_cap<F>(v: &mut Vec<u8>, f: F) -> io::Result<usize> |  | ||||||
|     where F: FnOnce(&mut [u8]) -> io::Result<usize> |  | ||||||
| { |  | ||||||
|     let len = v.len(); |  | ||||||
|     let new_area = v.capacity() - len; |  | ||||||
|     v.extend(iter::repeat(0).take(new_area)); |  | ||||||
|     match f(&mut v[len..]) { |  | ||||||
|         Ok(n) => { |  | ||||||
|             v.truncate(len + n); |  | ||||||
|             Ok(n) |  | ||||||
|         } |  | ||||||
|         Err(e) => { |  | ||||||
|             v.truncate(len); |  | ||||||
|             Err(e) |  | ||||||
|         } |  | ||||||
|     } |  | ||||||
| } |  | ||||||
|  |  | ||||||
| #[inline] |  | ||||||
| fn reserve(v: &mut Vec<u8>) { |  | ||||||
|     let cap = v.capacity(); |  | ||||||
|     if v.len() == cap { |  | ||||||
|         v.reserve(cmp::min(cap * 4, MAX_BUFFER_SIZE) - cap); |  | ||||||
|     } |     } | ||||||
| } | } | ||||||
|  |  | ||||||
| #[cfg(test)] | #[cfg(test)] | ||||||
| mod tests { | mod tests { | ||||||
|  |  | ||||||
|     use std::io::BufRead; |     use std::io::{self, Read, BufRead}; | ||||||
|     use super::BufReader; |     use super::BufReader; | ||||||
|  |  | ||||||
|  |     struct SlowRead(u8); | ||||||
|  |  | ||||||
|  |     impl Read for SlowRead { | ||||||
|  |         fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> { | ||||||
|  |             let state = self.0; | ||||||
|  |             self.0 += 1; | ||||||
|  |             (&match state % 3 { | ||||||
|  |                 0 => b"foo", | ||||||
|  |                 1 => b"bar", | ||||||
|  |                 _ => b"baz", | ||||||
|  |             }[..]).read(buf) | ||||||
|  |         } | ||||||
|  |     } | ||||||
|  |  | ||||||
|     #[test] |     #[test] | ||||||
|     fn test_consume_and_get_buf() { |     fn test_consume_and_get_buf() { | ||||||
|         let mut rdr = BufReader::new(&b"foo bar baz"[..]); |         let mut rdr = BufReader::new(SlowRead(0)); | ||||||
|         rdr.read_into_buf().unwrap(); |         rdr.read_into_buf().unwrap(); | ||||||
|         rdr.consume(8); |         rdr.consume(1); | ||||||
|  |         assert_eq!(rdr.get_buf(), b"oo"); | ||||||
|  |         rdr.read_into_buf().unwrap(); | ||||||
|  |         rdr.read_into_buf().unwrap(); | ||||||
|  |         assert_eq!(rdr.get_buf(), b"oobarbaz"); | ||||||
|  |         rdr.consume(5); | ||||||
|         assert_eq!(rdr.get_buf(), b"baz"); |         assert_eq!(rdr.get_buf(), b"baz"); | ||||||
|  |         rdr.consume(3); | ||||||
|  |         assert_eq!(rdr.get_buf(), b""); | ||||||
|  |         assert_eq!(rdr.pos, 0); | ||||||
|  |         assert_eq!(rdr.cap, 0); | ||||||
|     } |     } | ||||||
| } | } | ||||||
|   | |||||||
		Reference in New Issue
	
	Block a user