1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137

use std::io::{Error, Read};
use std::mem::swap;

/// from https://github.com/C4K3/peekable-reader-rs/blob/master/src/lib.rs

/// A wrapper around any struct implementing the `Read` trait, additionally
/// allowing for `peek` operations to be performed. Therefore, the
/// `PeekableReader` struct also implements the `Read` trait.
///
/// The primary invariant of the `PeekableReader` is that after calling the
/// `peek` method, the next `read_byte` call will return the same result as
/// the `peek` does. When the result is a byte (read off the wrapped reader),
/// any read-type method of the `Reader` trait will include the byte as the
/// first one. On the other hand, if the result is an `io::Error`, the error
/// will be returned regardless of which read-type method of the `Reader` is
/// invoked. Consecutive `peek`s before any read-type operation is used
/// always return the same `io::Result`.
///
/// When using peek_u8 you will get a Result<u8, &Error>, while when you call
/// any of the `Read` functions you will get an io::Result<u8>, thereby
/// consuming the io::Error (if any.)
pub struct PeekableReader<R> {
    inner: R,
    peeked_result: Option<Result<u8, Error>>,
}

/// A wrapper around any struct implementing the `Read` trait, additionally
/// allowing for `peek` operations to be performed. Therefore, the
/// `PeekableReader` struct also implements the `Read` trait.
///
/// The primary invariant of the `PeekableReader` is that after calling the
/// `peek` method, the next `read_byte` call will return the same result as
/// the `peek` does. When the result is a byte (read off the wrapped reader),
/// any read-type method of the `Reader` trait will include the byte as the
/// first one. On the other hand, if the result is an `io::Error`, the error
/// will be returned regardless of which read-type method of the `Reader` is
/// invoked. Consecutive `peek`s before any read-type operation is used
/// always return the same `io::Result`.
///
pub trait Peekable: Read {
    /// When using peek_u8 you will get a Result<u8, &Error>, while when you call
    /// any of the `Read` functions you will get an io::Result<u8>, thereby
    /// consuming the io::Error (if any.)
    fn peek_u8(&mut self) -> Result<u8, &Error>;
}

impl<R: Read> Peekable for PeekableReader<R> {
    /// Returns the `io::Result` which the Reader will return on the next
    /// `get_byte` call.
    ///
    /// If the `io::Result` is indeed an `io::Error`, the error will be returned
    /// for any subsequent read operation invoked upon the `Read`er.
    fn peek_u8(&mut self) -> Result<u8, &Error> {
        // Return either the currently cached peeked byte or obtain a new one
        // from the underlying reader.
        match self.peeked_result {
            Some(Ok(x)) => Ok(x),
            Some(Err(ref e)) => Err(e),
            None => {
                // First get the result of the read from the underlying reader
                let mut tmp: [u8; 1] = [0];
                self.peeked_result = match self.inner.read_exact(&mut tmp) {
                    Ok(()) => Some(Ok(tmp[0])),
                    Err(e) => Some(Err(e)),
                };

                // Now just return that
                let tmp: Result<u8, &Error> = match self.peeked_result {
                    Some(Ok(x)) => Ok(x),
                    Some(Err(ref e)) => Err(e),
                    None => unreachable!(),
                };
                tmp
            }
        }
    }
}

impl<R: Read> PeekableReader<R> {
    /// Initializes a new `PeekableReader` which wraps the given underlying
    /// reader.
    pub fn new(reader: R) -> PeekableReader<R> {
        PeekableReader {
            inner: reader,
            peeked_result: None,
        }
    }
}

impl<R: Read> Read for PeekableReader<R> {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize, Error> {
        if buf.is_empty() {
            return Ok(0);
        }

        // First, put the byte that was read off the underlying reader in a
        // (possible) previous peek operation (if such a byte is indeed cached)
        let mut tmp = None;
        swap(&mut tmp, &mut self.peeked_result);
        let offset = match tmp {
            Some(Err(e)) => {
                return Err(e);
            }
            Some(Ok(b)) => {
                buf[0] = b;
                1
            }
            None => 0,
        };

        if offset == 1 && buf.len() == 1 {
            // We've filled the buffer by using the previously peeked byte
            Ok(1)
        } else {
            // We are still missing more bytes, so we read them from the
            // underlying reader and place them directly in the correct place
            // in the buffer.
            Ok((self.inner.read(&mut buf[offset..]))? + offset)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::vlq_encode::ReadSigmaVlqExt;
    use std::io::Cursor;

    #[test]
    fn test_peek_u8() {
        let mut r = PeekableReader::new(Cursor::new(vec![0, 1]));
        assert_eq!(r.peek_u8().unwrap(), 0);
        assert_eq!(r.get_u8().unwrap(), 0);
        assert_eq!(r.peek_u8().unwrap(), 1);
        assert_eq!(r.get_u8().unwrap(), 1);
    }
}