From 1c0d2b81b1ca34ad07c8193c4183e73340528db8 Mon Sep 17 00:00:00 2001 From: Raphael Speyer Date: Sat, 25 Apr 2015 08:04:18 +1000 Subject: [PATCH] Add IterReader --- src/libstd/io/iterator.rs | 213 ++++++++++++++++++++++++++++++++++++++ src/libstd/io/mod.rs | 2 + 2 files changed, 215 insertions(+) create mode 100644 src/libstd/io/iterator.rs diff --git a/src/libstd/io/iterator.rs b/src/libstd/io/iterator.rs new file mode 100644 index 0000000000000..d1d6bb9d58321 --- /dev/null +++ b/src/libstd/io/iterator.rs @@ -0,0 +1,213 @@ +#![crate_type = "lib"] + +use prelude::v1::*; +use io::prelude::*; + +use cmp; +use io::{self}; +use ptr; +use slice; +use sync; + +/// An `IterReader` implements std::io::Read and std::io::BufRead for an iterator of byte sequences. +pub struct IterReader where V : AsRef<[u8]>, I : Iterator { + iter: I, // The underlying source of the bytes we will read. + buf: Option, // Contains the last item that we got from the iterator. + pos: usize, // The number of bytes that have already been read from buf. + closed: bool, +} + +impl , I : Iterator> IterReader { + pub fn new(iter: I) -> IterReader { + IterReader { iter: iter, buf: None, pos: 0, closed: false, } + } + + fn buffered(&self) -> &[u8] { + match self.buf { + Some(ref buf) => buf.as_ref()[self.pos..].as_ref(), + None => &[], + } + } + + fn needs_refill(&self) -> bool { + !self.closed && self.buffered().len() == 0 + } + + fn refill(&mut self) { + self.pos = 0; + self.buf = self.iter.next(); + self.closed = self.buf.is_none() + } +} + +impl , I : Iterator> io::BufRead for IterReader { + fn fill_buf(&mut self) -> io::Result<&[u8]> { + while self.needs_refill() { // We loop in case the last item was empty. + self.refill(); + } + Ok(self.buffered()) + } + + fn consume(&mut self, amt: usize) { self.pos += amt; } +} + +impl , I : Iterator> io::Read for IterReader { + fn read(&mut self, buf: &mut [u8]) -> io::Result { + if self.needs_refill() { + self.refill(); + } + let amt = cmp::min(self.buffered().len(), buf.len()); + slice::bytes::copy_memory(&self.buffered()[..amt], &mut buf[..amt]); + self.consume(amt); + Ok(amt) + } +} + +#[test] +fn test_read_vec() { + use io::Read; + + let pieces = vec![ + vec![1, 2], + vec![], + vec![3, 4, 5], + vec![6, 7, 8, 9], + vec![10], + ]; + + let mut reader = IterReader::new(pieces.into_iter()); + + let mut buf = [0; 3]; + + assert_eq!(0, reader.read(&mut []).unwrap()); + + assert_eq!(2, reader.read(&mut buf).unwrap()); + assert_eq!([1,2,0], buf); + + assert_eq!(0, reader.read(&mut buf).unwrap()); + assert_eq!([1,2,0], buf); + + assert_eq!(3, reader.read(&mut buf).unwrap()); + assert_eq!([3,4,5], buf); + + assert_eq!(3, reader.read(&mut buf).unwrap()); + assert_eq!([6,7,8], buf); + + assert_eq!(1, reader.read(&mut buf).unwrap()); + assert_eq!([9,7,8], buf); + + assert_eq!(1, reader.read(&mut buf).unwrap()); + assert_eq!([10,7,8], buf); + + assert_eq!(0, reader.read(&mut buf).unwrap()); + assert_eq!(0, reader.read(&mut buf).unwrap()); +} + +#[test] +fn test_read_array() { + use io::Read; + + let pieces = vec![ + [1, 2], + [3, 4], + [5, 6], + [7, 8], + [9, 10], + ]; + + let mut reader = IterReader::new(pieces.into_iter()); + + let mut buf = [0; 3]; + + assert_eq!(0, reader.read(&mut []).unwrap()); + + assert_eq!(2, reader.read(&mut buf).unwrap()); + assert_eq!([1,2,0], buf); + + assert_eq!(0, reader.read(&mut buf).unwrap()); + assert_eq!([1,2,0], buf); + + assert_eq!(3, reader.read(&mut buf).unwrap()); + assert_eq!([3,4,5], buf); + + assert_eq!(3, reader.read(&mut buf).unwrap()); + assert_eq!([6,7,8], buf); + + assert_eq!(1, reader.read(&mut buf).unwrap()); + assert_eq!([9,7,8], buf); + + assert_eq!(1, reader.read(&mut buf).unwrap()); + assert_eq!([10,7,8], buf); + + assert_eq!(0, reader.read(&mut buf).unwrap()); + assert_eq!(0, reader.read(&mut buf).unwrap()); +} + +#[test] +fn test_read_from_chan() { + use io::Read; + + let mut reader = { + let (tx, rx) = sync::mpsc::channel(); + tx.send(vec![1, 2]).unwrap(); + tx.send(vec![]).unwrap(); + tx.send(vec![3, 4, 5]).unwrap(); + tx.send(vec![6, 7, 8, 9]).unwrap(); + tx.send(vec![10]).unwrap(); + + IterReader::new(rx.into_iter()) + }; + + let mut buf = [0; 3]; + + assert_eq!(0, reader.read(&mut []).unwrap()); + + assert_eq!(2, reader.read(&mut buf).unwrap()); + assert_eq!([1,2,0], buf); + + assert_eq!(0, reader.read(&mut buf).unwrap()); + assert_eq!([1,2,0], buf); + + assert_eq!(3, reader.read(&mut buf).unwrap()); + assert_eq!([3,4,5], buf); + + assert_eq!(3, reader.read(&mut buf).unwrap()); + assert_eq!([6,7,8], buf); + + assert_eq!(1, reader.read(&mut buf).unwrap()); + assert_eq!([9,7,8], buf); + + assert_eq!(1, reader.read(&mut buf).unwrap()); + assert_eq!([10,7,8], buf); + + assert_eq!(0, reader.read(&mut buf).unwrap()); + assert_eq!(0, reader.read(&mut buf).unwrap()); +} + +#[test] +fn test_bufread() { + let pieces = vec![ + b"he".to_vec(), + b"llo wo".to_vec(), + b"".to_vec(), + b"rld\n\nhow ".to_vec(), + b"are you?".to_vec(), + b"".to_vec(), + ]; + + let lines = vec![ + "hello world\n", + "\n", + "how are you?", + "", + "", + ]; + + let mut reader = IterReader::new(pieces.into_iter()); + + for line in lines { + let buf = &mut "".to_string(); + assert_eq!(line.len(), reader.read_line(buf).unwrap()); + assert_eq!(line.to_string(), *buf); + } +} diff --git a/src/libstd/io/mod.rs b/src/libstd/io/mod.rs index e8d15a95b806b..29dace31c5fd7 100644 --- a/src/libstd/io/mod.rs +++ b/src/libstd/io/mod.rs @@ -30,6 +30,7 @@ pub use self::buffered::{BufReader, BufWriter, BufStream, LineWriter}; pub use self::buffered::IntoInnerError; pub use self::cursor::Cursor; pub use self::error::{Result, Error, ErrorKind}; +pub use self::iterator::{IterReader}; pub use self::util::{copy, sink, Sink, empty, Empty, repeat, Repeat}; pub use self::stdio::{stdin, stdout, stderr, _print, Stdin, Stdout, Stderr}; pub use self::stdio::{StdoutLock, StderrLock, StdinLock}; @@ -43,6 +44,7 @@ mod buffered; mod cursor; mod error; mod impls; +mod iterator; mod util; mod stdio;