thermidor/therm_util/src/reader.rs

use std::{fs,io,iter,mem,slice,vec};
use bzip2::bufread::BzDecoder;

/// A `ByteReader` is just a tiny wrapper over a mutable byte iterator, so we
/// can parse things more easily.
pub struct ByteReader<Rd> {
    bytes: Rd,
}

impl<R: io::Read> ByteReader<iter::FilterMap<io::Bytes<R>, &'static Fn(io::Result<u8>) -> Option<u8>>>
{
    /// Create a ByteReader from any type that implement Read
    pub fn from_reader(r: R) -> Self {
        const MK_OK: &'static Fn(io::Result<u8>) -> Option<u8> = &io::Result::<u8>::ok;
        let bytes = r.bytes().filter_map(MK_OK);
        ByteReader { bytes: bytes }
    }
}

impl ByteReader<iter::FilterMap<io::Bytes<fs::File>,
                                &'static Fn(io::Result<u8>) -> Option<u8>>>
{
    /// Create a reader by opening a named file for reading
    pub fn from_file(path: &str) -> io::Result<Self> {
        let f = try!(fs::File::open(path));
        Ok(ByteReader::from_reader(f))
    }
}

impl ByteReader<vec::IntoIter<u8>> {
    /// Create a reader from a vector of u8s
    pub fn from_vec(lst: Vec<u8>) -> Self {
        ByteReader { bytes: lst.into_iter() }
    }
}

impl<'a> ByteReader<iter::Cloned<slice::Iter<'a, u8>>> {
    /// Create a reader from a borrowed slice, with a copy on each access
    pub fn from_slice(lst: &'a [u8]) -> Self {
        ByteReader { bytes: lst.iter().cloned() }
    }
}

impl<R: io::BufRead> ByteReader<iter::FilterMap<io::Bytes<BzDecoder<R>>, &'static Fn(io::Result<u8>) -> Option<u8>>> {
    pub fn from_compressed_reader(r: R) -> Self {
        ByteReader::from_reader(BzDecoder::new(r))
    }
}

impl ByteReader<iter::FilterMap<io::Bytes<BzDecoder<io::BufReader<fs::File>>>, &'static Fn(io::Result<u8>) -> Option<u8>>> {
    pub fn from_compressed_file(path: &str) -> io::Result<Self> {
        let f = try!(fs::File::open(path));
        Ok(ByteReader::from_compressed_reader(io::BufReader::new(f)))
    }
}

#[test]
fn reader_tests() {
    assert!(ByteReader::from_slice(&[0x00]).read_twip() == Ok(-32.0));
    assert!(ByteReader::from_slice(&[0x00]).read_prefix_int() == Ok(0x00));
    assert!(ByteReader::from_slice(&[0x7f]).read_prefix_int() == Ok(0x7f));
    assert!(ByteReader::from_slice(&[0x80,0xff]).read_prefix_int() == Ok(0xff));
    assert!(ByteReader::from_slice(&[0x03,0x00,0x00,0x00])
            .read_several(|r| r.read_twip()) ==
            Ok(vec![-32.0,-32.0,-32.0]));
}

impl<T> ByteReader<T> where T: Iterator<Item=u8> {
    /// This gets the next byte, or fails if it's out of input.
    pub fn next(&mut self) -> Result<u8, String> {
        Ok(try!(self.bytes.next().ok_or("out of input")))
    }

    /// This reads a one-byte or two-byte float in the rough range of
    /// (192 .. -128). We're going to treat all models as if they're
    /// supposed to be centered in a 64x64x64 square, so this gives
    /// us 128 units on either side of the central square, as well.
    pub fn read_twip(&mut self) -> Result<f32, String> {
        let b1 = try!(self.next());
        if (b1 & 0x80) != 0 {
            let b2 = try!(self.next());
            let val = ((b1 as u16 & 0x7f) << 8) | b2 as u16;
            Ok((val as f32 / 102.0) - 128.0)
        } else {
            Ok((b1 as f32) - 32.0)
        }
    }

    /// This reads a single byte and treats it as a ratio.
    pub fn read_ratio(&mut self) -> Result<f32, String> {
        let b = try!(self.next());
        Ok(b as f32 / 255.0)
    }

    pub fn read_u64be(&mut self) -> Result<u64, String> {
        let a = try!(self.next());
        let b = try!(self.next());
        let c = try!(self.next());
        let d = try!(self.next());
        let e = try!(self.next());
        let f = try!(self.next());
        let g = try!(self.next());
        let h = try!(self.next());
        let rs = [ h, g, f, e, d, c, b, a ];
        unsafe { Ok(mem::transmute::<[u8;8],u64>(rs)) }
    }

    pub fn read_u64le(&mut self) -> Result<u64, String> {
        let a = try!(self.next());
        let b = try!(self.next());
        let c = try!(self.next());
        let d = try!(self.next());
        let e = try!(self.next());
        let f = try!(self.next());
        let g = try!(self.next());
        let h = try!(self.next());
        let rs = [ a, b, c, d, e, f, g, h ];
        unsafe { Ok(mem::transmute::<[u8;8],u64>(rs)) }
    }

    pub fn read_u32be(&mut self) -> Result<u32, String> {
        let a = try!(self.next());
        let b = try!(self.next());
        let c = try!(self.next());
        let d = try!(self.next());
        let rs = [ d, c, b, a ];
        unsafe { Ok(mem::transmute::<[u8;4],u32>(rs)) }
    }

    pub fn read_u32le(&mut self) -> Result<u32, String> {
        let a = try!(self.next());
        let b = try!(self.next());
        let c = try!(self.next());
        let d = try!(self.next());
        let rs = [ a, b, c, d ];
        unsafe { Ok(mem::transmute::<[u8;4],u32>(rs)) }
    }

    pub fn read_u16be(&mut self) -> Result<u16, String> {
        let a = try!(self.next());
        let b = try!(self.next());
        let rs = [ b, a ];
        unsafe { Ok(mem::transmute::<[u8;2],u16>(rs)) }
    }

    pub fn read_u16le(&mut self) -> Result<u16, String> {
        let a = try!(self.next());
        let b = try!(self.next());
        let rs = [ a, b ];
        unsafe { Ok(mem::transmute::<[u8;2],u16>(rs)) }
    }

    /// This reads a 64-bit int with a packed PrefixInteger representation.
    /// The shorter the int, the shorter the representation.
    pub fn read_prefix_int(&mut self) -> Result<u64, String> {
        fn match_bits(n: u8, mask: u8) -> bool {
            n & mask == mask
        }
        let b = try!(self.next());
        if match_bits(b, 0xff)      { self.continue_prefix_int(8, 0) }
        else if match_bits(b, 0xfe) { self.continue_prefix_int(7, 0) }
        else if match_bits(b, 0xfc) { self.continue_prefix_int(6, b & 0x01) }
        else if match_bits(b, 0xf8) { self.continue_prefix_int(5, b & 0x03) }
        else if match_bits(b, 0xf0) { self.continue_prefix_int(4, b & 0x07) }
        else if match_bits(b, 0xe0) { self.continue_prefix_int(3, b & 0x0f) }
        else if match_bits(b, 0xc0) { self.continue_prefix_int(2, b & 0x1f) }
        else if match_bits(b, 0x80) { self.continue_prefix_int(1, b & 0x3f) }
        else {                        self.continue_prefix_int(0, b) }
    }

    /// This is a helper function for parsing prefix ints, too.
    fn continue_prefix_int(&mut self, mut left: u8, upper: u8) -> Result<u64, String> {
        let mut ret = upper as u64;
        while left > 0 {
            left -= 1;
            ret = (ret << 8) | try!(self.next()) as u64;
        }
        Ok(ret)
    }

    /// This reads a PrefixInteger to find out how many other things to read,
    /// and then reads that number of things.
    pub fn read_several<F, R>(&mut self, reader: F) -> Result<Vec<R>, String>
        where F: Fn(&mut ByteReader<T>) -> Result<R, String>
    {
        let ct = try!(self.read_prefix_int());
        let mut ret = Vec::with_capacity(ct as usize);
        for _ in 0..ct {
            ret.push(try!(reader(self)))
        }
        Ok(ret)
    }

    /// This reads a PrefixInteger number of bytes, and then parses those
    /// bytes as a UTF-8 string. This means, importantly, that we cannot
    /// naïvely produce values intended to be parsed with this using a
    /// basic string length.
    pub fn read_string(&mut self) -> Result<String, String> {
        let raw_bytes = try!(self.read_several(|r| r.next()));
        match String::from_utf8(raw_bytes) {
            Ok(s) => Ok(s),
            Err(e) => Err(format!("Exception when parsing UTF-8: {:?}", e)),
        }
    }
}