lz-string 0.1.1

use std::iter::Map;
use std::str::Bytes;

#[derive(Debug, PartialEq)]
pub enum Error {
    InvalidByte(u8),
}

pub struct DecodeIterator<Input> {
    bytes: Input,
    decoded: [u8; 3],
    decoded_idx: u8,
}

type MapFn = fn(u8) -> Result<u8, Error>;
type MappedBytes<'a> = Map<Bytes<'a>, MapFn>;

impl<'a> DecodeIterator<MappedBytes<'a>> {
    pub fn from_str(str: &'a str) -> Self {
        let bytes = str.bytes().map(map_byte as MapFn);
        DecodeIterator {
            bytes,
            decoded: [0, 0, 0],
            decoded_idx: 0,
        }
    }
}

impl<I> DecodeIterator<I> {
    fn has_decoded(&self) -> bool {
        self.decoded_idx != 0 && (self.decoded_idx as usize) < self.decoded.len()
    }
}

impl<I> DecodeIterator<I>
where
    I: Iterator<Item = Result<u8, Error>>,
{
    fn pull_chunk(&mut self) -> Option<Result<[u8; 4], Error>> {
        let mut res = [0, 0, 0, 0];

        res[0] = nested_try!(self.bytes.next()?);
        res[1] = nested_try!(self.bytes.next().unwrap_or(Ok(0)));
        res[2] = nested_try!(self.bytes.next().unwrap_or(Ok(0)));
        res[3] = nested_try!(self.bytes.next().unwrap_or(Ok(0)));

        Some(Ok(res))
    }

    fn decode_chunk(&mut self) -> Option<Result<(), Error>> {
        let chunk = nested_try!(self.pull_chunk()?);

        self.decoded[0] = (chunk[0] << 2) | ((chunk[1] & 0x30) >> 4);
        self.decoded[1] = ((chunk[1] & 0x0F) << 4) | ((chunk[2] & 0x3C) >> 2);
        self.decoded[2] = ((chunk[2] & 0x03) << 6) | chunk[3];

        self.decoded_idx = 0;

        Some(Ok(()))
    }
}

impl<I> Iterator for DecodeIterator<I>
where
    I: Iterator<Item = Result<u8, Error>>,
{
    type Item = Result<u8, Error>;

    fn next(&mut self) -> Option<Self::Item> {
        if !self.has_decoded() {
            nested_try!(self.decode_chunk()?);
        }

        self.decoded_idx += 1;
        Some(Ok(self.decoded[(self.decoded_idx - 1) as usize]))
    }
}

fn map_byte(b: u8) -> Result<u8, Error> {
    TABLE[b as usize].ok_or(Error::InvalidByte(b))
}

const TABLE: [Option<u8>; 256] = [
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    Some(62),
    None,
    None,
    None,
    Some(63),
    Some(52),
    Some(53),
    Some(54),
    Some(55),
    Some(56),
    Some(57),
    Some(58),
    Some(59),
    Some(60),
    Some(61),
    None,
    None,
    None,
    Some(0),
    None,
    None,
    None,
    Some(0),
    Some(1),
    Some(2),
    Some(3),
    Some(4),
    Some(5),
    Some(6),
    Some(7),
    Some(8),
    Some(9),
    Some(10),
    Some(11),
    Some(12),
    Some(13),
    Some(14),
    Some(15),
    Some(16),
    Some(17),
    Some(18),
    Some(19),
    Some(20),
    Some(21),
    Some(22),
    Some(23),
    Some(24),
    Some(25),
    None,
    None,
    None,
    None,
    None,
    None,
    Some(26),
    Some(27),
    Some(28),
    Some(29),
    Some(30),
    Some(31),
    Some(32),
    Some(33),
    Some(34),
    Some(35),
    Some(36),
    Some(37),
    Some(38),
    Some(39),
    Some(40),
    Some(41),
    Some(42),
    Some(43),
    Some(44),
    Some(45),
    Some(46),
    Some(47),
    Some(48),
    Some(49),
    Some(50),
    Some(51),
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
];

#[cfg(test)]
mod test {
    use super::{DecodeIterator, Error};

    #[test]
    fn empty() {
        let data = "";
        let mut iterator = DecodeIterator::from_str(data);

        assert_eq!(iterator.next(), None);
    }

    #[test]
    fn simple() {
        let data = "AAAA";
        let mut iterator = DecodeIterator::from_str(data);

        assert_eq!(iterator.next(), Some(Ok(0)));
        assert_eq!(iterator.next(), Some(Ok(0)));
        assert_eq!(iterator.next(), Some(Ok(0)));
        assert_eq!(iterator.next(), None);
    }

    #[test]
    fn mixed() {
        let data = "aE0/";
        let mut iterator = DecodeIterator::from_str(data);

        assert_eq!(iterator.next(), Some(Ok(0x68)));
        assert_eq!(iterator.next(), Some(Ok(0x4d)));
        assert_eq!(iterator.next(), Some(Ok(0x3f)));
        assert_eq!(iterator.next(), None);
    }

    #[test]
    fn incomplete() {
        let data = "b5";
        let mut iterator = DecodeIterator::from_str(data);

        assert_eq!(iterator.next(), Some(Ok(0x6f)));
        assert_eq!(iterator.next(), Some(Ok(0x90)));
        assert_eq!(iterator.next(), Some(Ok(0x00)));
        assert_eq!(iterator.next(), None);
    }

    #[test]
    fn multiple() {
        let data = "aE0/b5";
        let mut iterator = DecodeIterator::from_str(data);

        assert_eq!(iterator.next(), Some(Ok(0x68)));
        assert_eq!(iterator.next(), Some(Ok(0x4d)));
        assert_eq!(iterator.next(), Some(Ok(0x3f)));
        assert_eq!(iterator.next(), Some(Ok(0x6f)));
        assert_eq!(iterator.next(), Some(Ok(0x90)));
        assert_eq!(iterator.next(), Some(Ok(0x00)));
        assert_eq!(iterator.next(), None);
    }

    #[test]
    fn invalid() {
        let data = "*";
        let mut iterator = DecodeIterator::from_str(data);

        assert_eq!(iterator.next(), Some(Err(Error::InvalidByte(0x2A))));
        assert_eq!(iterator.next(), None);
    }
}