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: W, buffer: IOList, items: Vec>, streaming: bool, } impl PackedWriter<&mut Vec> { #[inline(always)] pub fn encode>( enc: &mut Enc, v: &N, ) -> io::Result> { let mut buf: Vec = Vec::new(); PackedWriter::new(&mut buf).write(enc, v)?; Ok(buf) } #[inline(always)] pub fn encode_iovalue(v: &IOValue) -> io::Result> { 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 PackedWriter { #[inline(always)] pub fn new(write: W) -> Self { PackedWriter { w: write, buffer: IOList::new(), items: Vec::with_capacity(16), streaming: false, } } #[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<()> { if self.buffer.len() > 0 { let buffer = std::mem::replace(&mut self.buffer, IOList::new()); match self.items.last_mut() { Some(iols) => iols.push(buffer), None => { if self.streaming { let mut len = IOList::new(); varint(&mut len, buffer.len()); len.write_to(&mut self.w)?; } buffer.write_to(&mut self.w)?; } } } Ok(()) } fn at_toplevel(&self) -> bool { self.items.is_empty() } pub fn finish_item_if_toplevel(&mut self) -> io::Result<()> { if self.at_toplevel() { self.finish_item()?; } Ok(()) } pub fn start_seq(&mut self) -> io::Result<()> { self.items.push(Vec::with_capacity(32)); 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 i in items { varint(&mut self.buffer, i.len()); self.buffer.append(i); } self.finish_item_if_toplevel() } pub fn streaming(mut self) -> Self { if !self.at_toplevel() { panic!("Attempt to enable streaming mode mid-way through serialization"); } self.streaming = true; self } } macro_rules! fits_in_bytes { ($v:ident, $limit:literal) => ({ let bits = $limit * 8 - 1; $v >= -(2 << bits) && $v < (2 << bits) }) } impl Writer for PackedWriter { #[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<()> { if self.at_toplevel() && !self.streaming { self.write_tag(Tag::Sequence); self.finish_item()?; self.streaming = true; Ok(()) } else { self.start_seq() } } #[inline(always)] fn end_sequence(&mut self) -> io::Result<()> { if self.at_toplevel() && self.streaming { self.streaming = false; Ok(()) } else { 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 (ki, 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>( &mut self, enc: &mut Enc, v: &N, ) -> io::Result<()> { self.inner_write(enc, v)?; self.finish_item_if_toplevel() } }