preserves/implementations/rust/preserves/src/value/packed/writer.rs

use num::bigint::BigInt;
use num::cast::ToPrimitive;
use std::borrow::Cow;
use std::convert::TryInto;
use std::io;
use super::constants::Tag;
use super::super::DomainEncode;
use super::super::IOValue;
use super::super::IOValueDomainCodec;
use super::super::NestedValue;
use super::super::boundary as B;
use super::super::iolist::IOList;
use super::super::writer::Writer;

pub struct PackedWriter<W: io::Write> {
    w: W,
    buffer: IOList,
    items: Vec<Vec<IOList>>,
}

impl PackedWriter<&mut Vec<u8>> {
    #[inline(always)]
    pub fn encode<N: NestedValue, Enc: DomainEncode<N::Embedded>>(
        enc: &mut Enc,
        v: &N,
    ) -> io::Result<Vec<u8>> {
        let mut buf: Vec<u8> = Vec::new();
        PackedWriter::new(&mut buf).write(enc, v)?;
        Ok(buf)
    }

    #[inline(always)]
    pub fn encode_iovalue(v: &IOValue) -> io::Result<Vec<u8>> {
        Self::encode(&mut IOValueDomainCodec, v)
    }
}

pub fn varint(iol: &mut IOList, v: usize) {
    _varint(iol, v, 0x80)
}

fn _varint(iol: &mut IOList, v: usize, d: u8) {
    if v < 128 {
        iol.write(v as u8 + d);
    } else {
        _varint(iol, v >> 7, 0);
        iol.write((v & 0x7f) as u8 + d);
    }
}

impl<W: io::Write> PackedWriter<W> {
    #[inline(always)]
    pub fn new(write: W) -> Self {
        PackedWriter {
            w: write,
            buffer: IOList::new(),
            items: Vec::new(),
        }
    }

    #[inline(always)]
    pub fn write_byte(&mut self, b: u8) {
        self.buffer.write(b)
    }

    #[inline(always)]
    pub fn write_all(&mut self, bs: &[u8]) {
        self.buffer.write_all(Cow::Borrowed(bs))
    }

    #[inline(always)]
    pub fn write_integer(&mut self, bs: &[u8]) -> io::Result<()> {
        self.write_atom(Tag::SignedInteger, bs)
    }

    #[inline(always)]
    pub fn write_tag(&mut self, tag: Tag) {
        self.write_byte(tag.into())
    }

    #[inline(always)]
    pub fn write_atom(&mut self, tag: Tag, bs: &[u8]) -> io::Result<()> {
        self.write_tag(tag);
        self.write_all(bs);
        self.finish_item_if_toplevel()
    }

    pub fn finish_item(&mut self) -> io::Result<()> {
        let buffer = std::mem::replace(&mut self.buffer, IOList::new());
        match self.items.last_mut() {
            Some(iols) => Ok(iols.push(buffer)),
            None => buffer.write_to(&mut self.w),
        }
    }

    pub fn finish_item_if_toplevel(&mut self) -> io::Result<()> {
        if self.items.is_empty() {
            self.finish_item()?;
        }
        Ok(())
    }

    pub fn start_seq(&mut self) -> io::Result<()> {
        self.items.push(Vec::new());
        Ok(())
    }

    pub fn finish_seq(&mut self, tag: Tag, sort: bool) -> io::Result<()> {
        let mut items = self.items.pop().unwrap();
        if sort {
            items.sort();
        }
        self.write_tag(tag);
        for mut i in items {
            varint(&mut self.buffer, i.len());
            self.buffer.append(i);
        }
        self.finish_item_if_toplevel()
    }
}

macro_rules! fits_in_bytes {
    ($v:ident, $limit:literal) => ({
        let bits = $limit * 8 - 1;
        $v >= -(2 << bits) && $v < (2 << bits)
    })
}

impl<W: io::Write> Writer for PackedWriter<W>
{
    #[inline(always)]
    fn boundary(&mut self, b: &B::Type) -> io::Result<()> {
        match (b.closing.as_ref(), b.opening.as_ref()) {
            (Some(_), _) =>
                self.finish_item(),
            (None, _) =>
                Ok(()),
        }
    }

    #[inline(always)]
    fn start_annotations(&mut self) -> io::Result<()> {
        self.start_seq()
    }

    #[inline(always)]
    fn end_annotations(&mut self) -> io::Result<()> {
        self.items.last_mut().unwrap().rotate_right(1);
        self.finish_seq(Tag::Annotation, false)
    }

    #[inline(always)]
    fn write_bool(&mut self, v: bool) -> io::Result<()> {
        self.write_tag(if v { Tag::True } else { Tag::False });
        self.finish_item_if_toplevel()
    }

    #[inline(always)]
    fn write_f32(&mut self, v: f32) -> io::Result<()> {
        self.write_tag(Tag::Float);
        self.write_all(&u32::to_be_bytes(f32::to_bits(v)));
        self.finish_item_if_toplevel()
    }

    #[inline(always)]
    fn write_f64(&mut self, v: f64) -> io::Result<()> {
        self.write_tag(Tag::Float);
        self.write_all(&u64::to_be_bytes(f64::to_bits(v)));
        self.finish_item_if_toplevel()
    }

    #[inline(always)]
    fn write_i8(&mut self, v: i8) -> io::Result<()> {
        if v == 0 {
            self.write_integer(&[])
        } else {
            self.write_integer(&[v as u8])
        }
    }

    #[inline(always)]
    fn write_u8(&mut self, v: u8) -> io::Result<()> {
        if let Ok(w) = v.try_into() { return self.write_i8(w) }
        self.write_integer(&[0, v])
    }

    #[inline(always)]
    fn write_i16(&mut self, v: i16) -> io::Result<()> {
        if let Ok(w) = v.try_into() { return self.write_i8(w) }
        self.write_integer(&[(v >> 8) as u8, (v & 255) as u8])
    }

    #[inline(always)]
    fn write_u16(&mut self, v: u16) -> io::Result<()> {
        if let Ok(w) = v.try_into() { return self.write_i16(w) }
        self.write_integer(&[0, (v >> 8) as u8, (v & 255) as u8])
    }

    #[inline(always)]
    fn write_i32(&mut self, v: i32) -> io::Result<()> {
        if let Ok(w) = v.try_into() { return self.write_i16(w) }
        if fits_in_bytes!(v, 3) {
            return self.write_integer(&[(v >> 16) as u8,
                                        (v >> 8) as u8,
                                        (v & 255) as u8]);
        }
        self.write_integer(&[(v >> 24) as u8,
                             (v >> 16) as u8,
                             (v >> 8) as u8,
                             (v & 255) as u8])
    }

    #[inline(always)]
    fn write_u32(&mut self, v: u32) -> io::Result<()> {
        if let Ok(w) = v.try_into() { return self.write_i32(w) }
        self.write_integer(&[0,
                             (v >> 24) as u8,
                             (v >> 16) as u8,
                             (v >> 8) as u8,
                             (v & 255) as u8])
    }

    #[inline(always)]
    fn write_i64(&mut self, v: i64) -> io::Result<()> {
        if let Ok(w) = v.try_into() { return self.write_i32(w) }
        if fits_in_bytes!(v, 5) {
            return self.write_integer(&[(v >> 32) as u8,
                                        (v >> 24) as u8,
                                        (v >> 16) as u8,
                                        (v >> 8) as u8,
                                        (v & 255) as u8]);
        }
        if fits_in_bytes!(v, 6) {
            return self.write_integer(&[(v >> 40) as u8,
                                        (v >> 32) as u8,
                                        (v >> 24) as u8,
                                        (v >> 16) as u8,
                                        (v >> 8) as u8,
                                        (v & 255) as u8]);
        }
        if fits_in_bytes!(v, 7) {
            return self.write_integer(&[(v >> 48) as u8,
                                        (v >> 40) as u8,
                                        (v >> 32) as u8,
                                        (v >> 24) as u8,
                                        (v >> 16) as u8,
                                        (v >> 8) as u8,
                                        (v & 255) as u8]);
        }
        self.write_integer(&[(v >> 56) as u8,
                             (v >> 48) as u8,
                             (v >> 40) as u8,
                             (v >> 32) as u8,
                             (v >> 24) as u8,
                             (v >> 16) as u8,
                             (v >> 8) as u8,
                             (v & 255) as u8])
    }

    #[inline(always)]
    fn write_u64(&mut self, v: u64) -> io::Result<()> {
        if let Ok(w) = v.try_into() { return self.write_i64(w) }
        self.write_integer(&[0,
                             (v >> 56) as u8,
                             (v >> 48) as u8,
                             (v >> 40) as u8,
                             (v >> 32) as u8,
                             (v >> 24) as u8,
                             (v >> 16) as u8,
                             (v >> 8) as u8,
                             (v & 255) as u8])
    }

    #[inline(always)]
    fn write_i128(&mut self, v: i128) -> io::Result<()> {
        if let Ok(w) = v.try_into() { return self.write_i64(w) }
        let bs: [u8; 16] = v.to_be_bytes();
        if fits_in_bytes!(v, 9)  { return self.write_integer(&bs[7..]); }
        if fits_in_bytes!(v, 10) { return self.write_integer(&bs[6..]); }
        if fits_in_bytes!(v, 11) { return self.write_integer(&bs[5..]); }
        if fits_in_bytes!(v, 12) { return self.write_integer(&bs[4..]); }
        if fits_in_bytes!(v, 13) { return self.write_integer(&bs[3..]); }
        if fits_in_bytes!(v, 14) { return self.write_integer(&bs[2..]); }
        if fits_in_bytes!(v, 15) { return self.write_integer(&bs[1..]); }
        self.write_integer(&bs)
    }

    #[inline(always)]
    fn write_u128(&mut self, v: u128) -> io::Result<()> {
        if let Ok(w) = v.try_into() { return self.write_i128(w) }
        let bs: [u8; 16] = v.to_be_bytes();
        self.write_tag(Tag::SignedInteger);
        self.write_byte(0);
        self.write_all(&bs);
        Ok(())
    }

    #[inline(always)]
    fn write_int(&mut self, v: &BigInt) -> io::Result<()> {
        match v.to_i8() {
            Some(n) => self.write_i8(n),
            None => {
                match v.to_i128() {
                    Some(n) => self.write_i128(n),
                    None => self.write_atom(Tag::SignedInteger, &v.to_signed_bytes_be()),
                }
            }
        }
    }

    #[inline(always)]
    fn write_string(&mut self, v: &str) -> io::Result<()> {
        self.write_tag(Tag::String);
        self.write_all(v.as_bytes());
        self.write_byte(0);
        self.finish_item_if_toplevel()
    }

    #[inline(always)]
    fn write_bytes(&mut self, v: &[u8]) -> io::Result<()> {
        self.write_atom(Tag::ByteString, v)
    }

    #[inline(always)]
    fn write_symbol(&mut self, v: &str) -> io::Result<()> {
        self.write_atom(Tag::Symbol, v.as_bytes())
    }

    #[inline(always)]
    fn start_record(&mut self) -> io::Result<()> {
        self.start_seq()
    }

    #[inline(always)]
    fn end_record(&mut self) -> io::Result<()> {
        self.finish_seq(Tag::Record, false)
    }

    #[inline(always)]
    fn start_sequence(&mut self) -> io::Result<()> {
        self.start_seq()
    }

    #[inline(always)]
    fn end_sequence(&mut self) -> io::Result<()> {
        self.finish_seq(Tag::Sequence, false)
    }

    #[inline(always)]
    fn start_set(&mut self) -> io::Result<()> {
        self.start_seq()
    }

    #[inline(always)]
    fn end_set(&mut self) -> io::Result<()> {
        self.finish_seq(Tag::Set, true)
    }

    #[inline(always)]
    fn start_dictionary(&mut self) -> io::Result<()> {
        self.start_seq()
    }

    #[inline(always)]
    fn end_dictionary(&mut self) -> io::Result<()> {
        let mut items_iter = self.items.pop().unwrap().into_iter();
        let mut chunks = vec![];
        while let Some(ki) = items_iter.next() {
            match items_iter.next() {
                Some(vi) => chunks.push((ki, vi)),
                None => panic!("Missing dictionary value during serialization"),
            }
        }
        chunks.sort();
        self.write_tag(Tag::Dictionary);
        for (mut ki, mut vi) in chunks {
            varint(&mut self.buffer, ki.len());
            self.buffer.append(ki);
            varint(&mut self.buffer, vi.len());
            self.buffer.append(vi);
        }
        Ok(())
    }

    #[inline(always)]
    fn start_embedded(&mut self) -> io::Result<()> {
        self.write_tag(Tag::Embedded);
        Ok(())
    }

    #[inline(always)]
    fn end_embedded(&mut self) -> io::Result<()> {
        Ok(())
    }

    #[inline(always)]
    fn flush(&mut self) -> io::Result<()> {
        if self.buffer.len() != 0 { panic!("Attempt to flush with unfinished item in buffer"); }
        self.w.flush()
    }

    fn write<N: NestedValue, Enc: DomainEncode<N::Embedded>>(
        &mut self,
        enc: &mut Enc,
        v: &N,
    ) -> io::Result<()> {
        self.inner_write(enc, v)?;
        self.finish_item_if_toplevel()
    }
}