preserves/implementations/rust/oo/src/packed/view.rs

use std::borrow::Cow;
use std::io;
use std::marker::PhantomData;

use num_bigint::BigInt;

use crate::IOValue;
use crate::PlainValue;
use crate::SignedInteger;
use crate::Value;
use crate::ValueClass;
use crate::ValueImpl;
use crate::error;
use crate::error::io_eof;
use crate::iovalue;
use crate::owned;
use crate::shell;
use crate::reader::NextToken;

use super::constants::Tag;

#[derive(Debug, Clone)]
pub struct View<'de, Packed: AsRef<[u8]> + 'de> {
    packed: Packed,
    value_offset: usize, // annotation_offset is implicitly 0
    value_end: usize,
    phantom: PhantomData<&'de ()>,
}

#[inline(always)]
fn at(packed: &[u8], i: usize) -> io::Result<u8> {
    match packed.get(i) {
        Some(v) => Ok(*v),
        None => Err(io_eof()),
    }
}

#[inline(always)]
fn tag_at(packed: &[u8], i: usize) -> io::Result<Tag> {
    Ok(Tag::try_from(at(packed, i)?)?)
}

fn skip_annotations(packed: &[u8], mut i: usize) -> io::Result<usize> {
    loop {
        if tag_at(packed, i)? == Tag::Annotation {
            i = skip_value(packed, i + 1)?;
        } else {
            return Ok(i);
        }
    }
}

fn varint(packed: &[u8], mut i: usize) -> io::Result<(u64, usize)> {
    let mut shift = 0;
    let mut acc: u64 = 0;
    loop {
        let v = at(packed, i)?;
        i = i + 1;
        if shift == 63 && v > 1 { Err(error::Error::Message("PackedReader length too long".to_string()))? }
        acc |= ((v & 0x7f) as u64) << shift;
        shift += 7;
        if v & 0x80 == 0 { return Ok((acc, i)) }
        if shift >= 70 { Err(error::Error::Message("PackedReader length too long".to_string()))? }
    }
}

fn skip_value(packed: &[u8], mut i: usize) -> io::Result<usize> {
    loop {
        let next_i = match tag_at(packed, i)? {
            Tag::False => i + 1,
            Tag::True => i + 1,
            Tag::Float => i + 5,
            Tag::Double => i + 9,
            Tag::End => Err(io::Error::new(io::ErrorKind::InvalidData, "Unexpected end tag"))?,
            Tag::Annotation => {
                i = skip_value(packed, i + 1)?;
                continue;
            }
            Tag::Embedded => {
                i = i + 1;
                continue;
            }
            Tag::SmallInteger(_) => i + 1,
            Tag::MediumInteger(n) => i + 1 + (n as usize),
            Tag::SignedInteger | Tag::ByteString => {
                let (n, i) = varint(packed, i + 1)?;
                i + (n as usize)
            }
            Tag::String | Tag::Symbol => {
                let (n, i) = varint(packed, i + 1)?;
                std::str::from_utf8(&packed[i .. i + (n as usize)]).map_err(
                    |_| io::Error::new(io::ErrorKind::InvalidData, "Invalid UTF-8"))?;
                i + (n as usize)
            }
            Tag::Record | Tag::Sequence | Tag::Set | Tag::Dictionary => {
                i = i + 1;
                while tag_at(packed, i)? != Tag::End {
                    i = skip_value(packed, i)?;
                }
                i + 1
            }
        };
        return Ok(next_i);
    }
}

impl<'de, Packed: AsRef<[u8]> + 'de> View<'de, Packed> {
    pub fn new(packed: Packed) -> io::Result<Self> {
        // println!("packed {:?}", &packed.as_ref());
        let value_offset = skip_annotations(packed.as_ref(), 0)?;
        let value_end = skip_value(packed.as_ref(), value_offset)?;
        if value_end > packed.as_ref().len() { Err(io_eof())? }
        Ok(View { packed, value_offset, value_end, phantom: PhantomData })
    }

    #[inline(always)]
    fn tag(&self) -> Tag {
        tag_at(self.packed.as_ref(), self.value_offset).unwrap()
    }

    #[inline(always)]
    fn varint(&self) -> (u64, usize) {
        varint(self.packed.as_ref(), self.value_offset + 1).unwrap()
    }

    #[inline(always)]
    fn sub(&self, offset: usize, len: usize) -> &[u8] {
        let packed = self.packed.as_ref();
        &packed[offset .. offset + len]
    }

    #[inline(always)]
    fn atom_chunk(&self) -> &[u8] {
        let (n, i) = self.varint();
        self.sub(i, n as usize)
    }

    #[inline(always)]
    fn signed_integer(&self, offset: usize, len: usize) -> SignedInteger {
        let bs = self.sub(offset, len);
        if (bs[0] & 0x80) == 0 {
            // Positive or zero.
            let mut i = 0;
            while i < len && bs[i] == 0 { i += 1; }
            if len - i <= 16 {
                let mut v: u128 = 0;
                for b in &bs[i..] { v = v << 8 | (*b as u128); }
                SignedInteger::from(v)
            } else {
                SignedInteger::from(Cow::Owned(BigInt::from_bytes_be(num_bigint::Sign::Plus, &bs[i..])))
            }
        } else {
            // Negative.
            let mut i = 0;
            while i < len && bs[i] == 0xff { i += 1; }
            if len - i <= 16 {
                let mut v: i128 = -1;
                for b in &bs[i..] { v = v << 8 | (*b as i128); }
                SignedInteger::from(v)
            } else {
                SignedInteger::from(Cow::Owned(BigInt::from_signed_bytes_be(&bs)))
            }
        }
    }
}

impl<'de, Packed: AsRef<[u8]> + 'de> ValueImpl<IOValue> for View<'de, Packed> {
    fn write(&self, w: &mut dyn crate::Writer, enc: &mut dyn crate::DomainEncode<IOValue>) -> io::Result<()> {
        crate::write_value(w, self, enc)
    }

    fn value_clone(&self) -> PlainValue<'static, IOValue> where IOValue: 'static {
        owned(View {
            packed: self.packed.as_ref()[..self.value_end].to_owned(),
            value_offset: self.value_offset,
            value_end: self.value_end,
            phantom: PhantomData,
        })
    }

    fn value_class(&self) -> ValueClass {
        match self.tag().into() {
            Some(NextToken::Annotation) | None => unreachable!(),
            Some(NextToken::Value(v)) => v,
        }
    }

    fn as_boolean(&self) -> Option<bool> {
        match self.tag() {
            Tag::False => Some(false),
            Tag::True => Some(true),
            _ => None,
        }
    }

    fn as_float(&self) -> Option<f32> {
        match self.tag() {
            Tag::Float => Some(f32::from_be_bytes(self.sub(self.value_offset + 1, 4).try_into().unwrap())),
            _ => None,
        }
    }

    fn as_double(&self) -> Option<f64> {
        match self.tag() {
            Tag::Double => Some(f64::from_be_bytes(self.sub(self.value_offset + 1, 8).try_into().unwrap())),
            _ => None,
        }
    }

    fn is_signed_integer(&self) -> bool {
        match self.tag() {
            Tag::SmallInteger(_) => true,
            Tag::MediumInteger(_) => true,
            Tag::SignedInteger => true,
            _ => false,
        }
    }

    fn as_signed_integer(&self) -> Option<SignedInteger> {
        match self.tag() {
            Tag::SmallInteger(v) => Some(v.into()),
            Tag::MediumInteger(n) => Some(self.signed_integer(self.value_offset + 1, n as usize)),
            Tag::SignedInteger => {
                let (n, i) = self.varint();
                Some(self.signed_integer(i, n as usize))
            }
            _ => None,
        }
    }

    fn as_string(&self) -> Option<Cow<'_, str>> {
        match self.tag() {
            Tag::String => Some(Cow::Borrowed(unsafe {
                // SAFETY: we already checked in the View constructor
                std::str::from_utf8_unchecked(self.atom_chunk())
            })),
            _ => None,
        }
    }

    fn as_bytestring(&self) -> Option<Cow<'_, [u8]>> {
        match self.tag() {
            Tag::ByteString => Some(Cow::Borrowed(self.atom_chunk())),
            _ => None,
        }
    }

    fn as_symbol(&self) -> Option<Cow<'_, str>> {
        match self.tag() {
            Tag::Symbol => Some(Cow::Borrowed(unsafe {
                // SAFETY: we already checked in the View constructor
                std::str::from_utf8_unchecked(self.atom_chunk())
            })),
            _ => None,
        }
    }

    fn is_record(&self) -> bool {
        self.tag() == Tag::Record
    }

    fn label(&self) -> Value<'_, IOValue> {
        if !self.is_record() { panic!("Not a record") }
        shell(View::new(&self.packed.as_ref()[self.value_offset + 1 .. self.value_end]).unwrap())
    }

    fn is_sequence(&self) -> bool {
        self.tag() == Tag::Sequence
    }

    fn len(&self) -> usize {
        self.iter().count()
    }

    fn index(&self, i: usize) -> Value<'_, IOValue> {
        self.iter().nth(i).unwrap()
    }

    fn iter(&self) -> Box<dyn Iterator<Item = Value<'_, IOValue>> + '_> {
        let mut i = Box::new(ViewIterator::new(&self.packed.as_ref()[self.value_offset + 1 .. self.value_end]));
        match self.tag() {
            Tag::Record => { i.next(); () }
            Tag::Sequence => (),
            Tag::Set => (),
            _ => panic!("Not iterable"),
        }
        i
    }

    fn is_set(&self) -> bool {
        self.tag() == Tag::Set
    }

    fn has(&self, v: &dyn ValueImpl<IOValue>) -> bool {
        self.iter().find(|e| v == &**e).is_some()
    }

    fn is_dictionary(&self) -> bool {
        self.tag() == Tag::Dictionary
    }

    fn get(&self, k: &dyn ValueImpl<IOValue>) -> Option<Value<'_, IOValue>> {
        for (kk, vv) in self.entries() {
            if &*kk == k { return Some(vv); }
        }
        None
    }

    fn entries(&self) -> Box<dyn Iterator<Item = (Value<'_, IOValue>, Value<'_, IOValue>)> + '_> {
        if !self.is_dictionary() { panic!("Not a dictionary") }
        Box::new(DictionaryAdapter(ViewIterator::new(&self.packed.as_ref()[self.value_offset + 1 .. self.value_end])))
    }

    fn is_embedded(&self) -> bool {
        self.tag() == Tag::Embedded
    }

    fn embedded(&self) -> Cow<'_, IOValue> {
        let bs = self.packed.as_ref()[self.value_offset + 1 .. self.value_end].to_vec();
        Cow::Owned(iovalue(View::new(bs).unwrap()))
    }

    fn annotations(&self) -> Option<Cow<'_, [IOValue]>> {
        if self.value_offset == 0 {
            None
        } else {
            let anns: Vec<IOValue> = AnnotationAdapter(ViewIterator::new(&self.packed.as_ref()[.. self.value_offset]))
                .map(|ann| iovalue(ann.into_owned()))
                .collect();
            Some(Cow::Owned(anns))
        }
    }
}

pub struct ViewIterator<'de> {
    packed: &'de [u8],
    offset: usize,
}

impl<'de> ViewIterator<'de> {
    pub fn new(packed: &'de [u8]) -> Self {
        ViewIterator { packed, offset: 0 }
    }
}

impl<'de> Iterator for ViewIterator<'de> {
    type Item = Value<'de, IOValue>;

    fn next(&mut self) -> Option<Self::Item> {
        let v = View::new(&self.packed[self.offset..]).ok()?;
        if v.tag() == Tag::End { return None; }
        self.offset += v.value_end;
        Some(shell(v))
    }
}

pub struct DictionaryAdapter<'de>(pub ViewIterator<'de>);

impl<'de> Iterator for DictionaryAdapter<'de> {
    type Item = (Value<'de, IOValue>, Value<'de, IOValue>);

    fn next(&mut self) -> Option<Self::Item> {
        let k = self.0.next()?;
        let v = self.0.next()?;
        Some((k, v))
    }
}

pub struct AnnotationAdapter<'de>(pub ViewIterator<'de>);

impl<'de> Iterator for AnnotationAdapter<'de> {
    type Item = Value<'de, IOValue>;

    fn next(&mut self) -> Option<Self::Item> {
        if let Ok(Tag::Annotation) = tag_at(self.0.packed.as_ref(), self.0.offset) {
            self.0.offset += 1;
            self.0.next()
        } else {
            None
        }
    }
}