openvariant 0.1.0

use std::{borrow::Cow, sync::Arc};

use crate::{
    grammar::{
        DictionaryLayout, OffsetWidth, ShortString, Variant, VariantArray, VariantHeader,
        VariantMetadata, VariantObject,
    },
    impl_read_for_datatype, impl_read_signed_leb128, impl_read_u24,
    validation::{is_strictly_increasing, parse_utf8},
};
use anyhow::{Context, Result, anyhow, bail, ensure};
use indexmap::IndexSet;

#[derive(Debug)]
pub struct VariantMetadataDecoder<'a> {
    cursor: usize,
    data: &'a [u8],
    header: VariantHeader,

    validate: bool,
}

impl<'a> VariantMetadataDecoder<'a> {
    pub fn try_new(metadata_buffer: &'a [u8], validate: bool) -> Result<Self> {
        let header = VariantHeader::try_from(
            *metadata_buffer
                .first()
                .ok_or_else(|| anyhow!("no header byte"))?,
        )?;

        Ok(Self {
            data: metadata_buffer,
            cursor: 1,
            header,
            validate,
        })
    }

    pub fn read_metadata(&mut self) -> Result<VariantMetadata<'a>> {
        let dictionary_size = self.read_offset()?;

        let dictionary_offsets = (0..dictionary_size
            .checked_add(1)
            .ok_or_else(|| anyhow!("overflow"))?)
            .map(|_| self.read_offset())
            .collect::<Result<Vec<_>>>()?;

        let bytes_len = self.data.len() - self.cursor;
        let bytes = self.read_slice(bytes_len as u32)?;

        let mut bytes_str = None;

        if self.validate {
            // Validation
            // 1. the first offset is 0
            ensure!(
                *dictionary_offsets
                    .first()
                    .expect("dictionary offsets must be nonempty")
                    == 0,
                "first offset must be 0"
            );

            // 2. the last offset will always be the total length of bytes
            ensure!(
                *dictionary_offsets
                    .last()
                    .expect("dictionary offsets must be nonempty") as usize
                    == bytes_len,
                "last offset must be the total length of bytes"
            );

            match self.header.layout {
                DictionaryLayout::SortedUnique => {
                    // strings in the dictionary must be unique and sorted in lexicographic order
                    let mut seen_strings = IndexSet::new();

                    let mut dictionary_offsets_iter = dictionary_offsets.windows(2);
                    while let Some([start, end]) = dictionary_offsets_iter.next() {
                        // 3. offsets are monotonically increasing
                        let slice = bytes
                            .get(*start as usize..*end as usize)
                            .ok_or_else(|| anyhow!("oob"))?;

                        // 4. the `bytes` array must be UTF-8 encoded strings
                        let str = str::from_utf8(slice)?;

                        ensure!(
                            seen_strings.insert(str),
                            "strings in dictionary must be unique"
                        );
                    }

                    ensure!(
                        is_strictly_increasing(seen_strings.iter()),
                        "strings in dictionary must be sorted in lexicographic order"
                    );
                }
                DictionaryLayout::Naive => {
                    // 3. offsets are non-decreasing
                    ensure!(
                        dictionary_offsets.is_sorted(),
                        "offsets must be non-decreasing"
                    );

                    // 4. the `bytes` array must be UTF-8 encoded strings
                    bytes_str =
                        Some(parse_utf8(bytes).context("bytes array must be utf8 encoded")?);
                }
            }
        }

        let bytes_str = if let Some(bytes_str) = bytes_str {
            bytes_str
        } else {
            parse_utf8(bytes)?
        };

        let mut strings = Vec::with_capacity(dictionary_size as usize);

        let mut dictionary_offsets_iter = dictionary_offsets.windows(2);
        while let Some([start, end]) = dictionary_offsets_iter.next() {
            strings.push(
                bytes_str
                    .get(*start as usize..*end as usize)
                    .expect("should never happen if we do validation or trust the input"),
            );
        }

        Ok(VariantMetadata {
            header: self.header,
            offsets: dictionary_offsets,
            strings,
        })
    }

    fn read_offset(&mut self) -> Result<u32> {
        match self.header.offset_width {
            OffsetWidth::U8 => self.read_u8().map(|b| b as u32),
            OffsetWidth::U16 => self.read_u16().map(|n| n as u32),
            OffsetWidth::U24 => self.read_u24(),
            OffsetWidth::U32 => self.read_u32(),
        }
    }

    impl_read_for_datatype!(read_u8, u8);
    impl_read_for_datatype!(read_u16, u16);
    impl_read_u24!();
    impl_read_for_datatype!(read_u32, u32);

    fn read_slice(&mut self, len: u32) -> Result<&'a [u8]> {
        let slice = self
            .data
            .get(self.cursor..self.cursor + len as usize)
            .ok_or_else(|| anyhow!("oob"))?;

        self.cursor += len as usize;

        Ok(slice)
    }
}

#[derive(Debug)]
pub struct VariantDecoder<'a> {
    variant_metadata: Arc<VariantMetadata<'a>>,

    value_metadata: u8,
    data: &'a [u8],
    cursor: usize,

    validate: bool,
}

impl<'a> VariantDecoder<'a> {
    pub fn try_new(metadata: &'a [u8], value: &'a [u8], validate: bool) -> Result<Self> {
        let variant_metadata =
            VariantMetadataDecoder::try_new(metadata, validate)?.read_metadata()?;

        Self::try_new_with_variant_metadata(Arc::new(variant_metadata), value, validate)
    }

    pub fn try_new_with_variant_metadata(
        variant_metadata: Arc<VariantMetadata<'a>>,
        value_buffer: &'a [u8],
        validate: bool,
    ) -> Result<Self> {
        let value_metadata = *value_buffer
            .first()
            .ok_or_else(|| anyhow!("expected metadata byte"))?;

        Ok(Self {
            variant_metadata,
            value_metadata,
            data: value_buffer,
            cursor: 1,
            validate,
        })
    }

    pub fn read_value(&mut self) -> Result<Variant<'a>> {
        let basic_type = self.value_metadata & 0b11;

        match basic_type {
            0 => self.read_primitive(),
            1 => self.read_short_string(),
            2 => self.read_object(),
            3 => self.read_array(),
            _ => unreachable!(),
        }
    }

    fn read_primitive(&mut self) -> Result<Variant<'a>> {
        let primitive_header = self.value_metadata >> 2;
        let v = match primitive_header {
            0 => Variant::Null,
            1 => Variant::BooleanTrue,
            2 => Variant::BooleanFalse,
            3 => Variant::Int8(self.read_i8()?),
            4 => Variant::Int16(self.read_i16()?),
            5 => Variant::Int32(self.read_i32()?),
            6 => Variant::Int64(self.read_i64()?),
            7 => Variant::Double(self.read_f64()?),
            8 => Variant::Decimal4(self.read_u8()?, self.read_i32()?),
            9 => Variant::Decimal8(self.read_u8()?, self.read_i64()?),
            10 => Variant::Decimal16(self.read_u8()?, self.read_i128()?),
            11 => Variant::Date(self.read_u32()?),
            12 => Variant::TimestampMicros(self.read_u64()?),
            13 => Variant::TimestampNTZMicros(self.read_u64()?),
            14 => Variant::Float(self.read_f32()?),
            15 => {
                let len = self.read_u32()?;
                Variant::Binary(Cow::Borrowed(self.read_slice(len)?))
            }
            16 => {
                let len = self.read_u32()?;
                let str = parse_utf8(self.read_slice(len)?)?;
                Variant::String(Cow::Borrowed(str))
            }
            17 => Variant::TimeNTZMicros(self.read_u64()?),
            18 => Variant::TimestampNanos(self.read_u64()?),
            19 => Variant::TimestampNTZNanos(self.read_u64()?),
            20 => Variant::UUID(self.read_u128()?),

            // experimental
            21 => Variant::Int16Leb128(self.read_i16_leb128()?),
            22 => Variant::Int32Leb128(self.read_i32_leb128()?),
            23 => Variant::Int64Leb128(self.read_i64_leb128()?),
            24 => Variant::Int128Leb128(self.read_i128_leb128()?),
            foreign => bail!("found unrecognized tag: {foreign}"),
        };

        Ok(v)
    }

    fn read_short_string(&mut self) -> Result<Variant<'a>> {
        // the short string header is the length of the string
        let short_string_header = self.value_metadata >> 2;
        let str = str::from_utf8(self.read_slice(short_string_header as u32)?)?;

        Ok(Variant::ShortString(ShortString::try_from(str)?))
    }

    fn read_object(&mut self) -> Result<Variant<'a>> {
        let object_header = self.value_metadata >> 2;

        let field_offset_width = OffsetWidth::try_from(object_header & 0b11)?;
        let field_id_width = OffsetWidth::try_from((object_header >> 2) & 0b11)?;
        let is_large = ((object_header >> 4) & 0b1) == 1;

        let num_elements = if is_large {
            self.read_u32()?
        } else {
            self.read_u8()? as u32
        };

        let field_ids = self.read_offsets(field_id_width, num_elements)?;
        let field_offsets = self.read_offsets(field_offset_width, num_elements + 1)?;

        if self.validate {
            {
                // objects can not have duplicate field names
                // since field ids always map to a field name, we should check for duplicate field ids
                let unique: IndexSet<u32> = IndexSet::from_iter(field_ids.iter().copied());
                ensure!(
                    unique.len() == field_ids.len(),
                    "Objects may not contain duplicate field names"
                );
            }

            match self.variant_metadata.header.layout {
                DictionaryLayout::SortedUnique => ensure!(
                    is_strictly_increasing(field_ids.iter()),
                    "field ids are not ordered by lexicographical order of the corresponding field names"
                ),
                DictionaryLayout::Naive => {
                    let is_field_ids_ordered_by_lexicographical_field_name = field_ids
                        .iter()
                        .map(|&field_id| self.variant_metadata.find_by_field_id(field_id))
                        .is_sorted_by(|a, b| match (a, b) {
                            (Some(a), Some(b)) => a <= b,
                            _ => false,
                        });

                    ensure!(
                        is_field_ids_ordered_by_lexicographical_field_name,
                        "field ids are not ordered by lexicographical order of the corresponding field names"
                    );
                }
            }
        }

        let mut variant_values = Vec::with_capacity(num_elements as usize);

        if let Some(last_offset) = field_offsets.last() {
            let value_buffer = self.read_slice(*last_offset)?;

            for &field_offset in field_offsets.iter().take(field_offsets.len() - 1) {
                let mut decoder = Self::try_new_with_variant_metadata(
                    Arc::clone(&self.variant_metadata),
                    value_buffer
                        .get(field_offset as usize..)
                        .ok_or_else(|| anyhow!("oob"))?,
                    self.validate,
                )?;

                variant_values.push(decoder.read_value()?);
            }
        }

        Ok(Variant::Object(VariantObject {
            field_ids,
            field_offsets,
            values: variant_values,
        }))
    }

    fn read_array(&mut self) -> Result<Variant<'a>> {
        let array_header = self.value_metadata >> 2;

        let field_offset_width = OffsetWidth::try_from(array_header & 0b11)?;
        let is_large = ((array_header >> 2) & 0b1) == 1;

        let num_elements = if is_large {
            self.read_u32()?
        } else {
            self.read_u8()? as u32
        };

        let field_offsets = self.read_offsets(field_offset_width, num_elements + 1)?;

        let mut variant_values = Vec::with_capacity(num_elements as usize);

        if let Some(last_offset) = field_offsets.last() {
            let value_buffer = self.read_slice(*last_offset)?;

            let mut field_offsets_iter = field_offsets.windows(2);
            while let Some([start, end]) = field_offsets_iter.next() {
                let value_buffer = value_buffer
                    .get(*start as usize..*end as usize)
                    .ok_or_else(|| anyhow!("oob"))?;

                let mut decoder = Self::try_new_with_variant_metadata(
                    Arc::clone(&self.variant_metadata),
                    value_buffer,
                    self.validate,
                )?;

                variant_values.push(decoder.read_value()?);
            }
        }

        Ok(Variant::Array(VariantArray {
            field_offsets,
            values: variant_values,
        }))
    }

    fn read_offset(&mut self, width: OffsetWidth) -> Result<u32> {
        match width {
            OffsetWidth::U8 => self.read_u8().map(|b| b as u32),
            OffsetWidth::U16 => self.read_u16().map(|n| n as u32),
            OffsetWidth::U24 => self.read_u24(),
            OffsetWidth::U32 => self.read_u32(),
        }
    }

    impl_read_for_datatype!(read_i8, i8);
    impl_read_for_datatype!(read_i16, i16);
    impl_read_for_datatype!(read_i32, i32);
    impl_read_for_datatype!(read_i64, i64);
    impl_read_for_datatype!(read_i128, i128);

    impl_read_for_datatype!(read_f32, f32);
    impl_read_for_datatype!(read_f64, f64);

    impl_read_for_datatype!(read_u8, u8);
    impl_read_for_datatype!(read_u16, u16);
    impl_read_u24!();
    impl_read_for_datatype!(read_u32, u32);
    impl_read_for_datatype!(read_u64, u64);
    impl_read_for_datatype!(read_u128, u128);

    fn read_slice(&mut self, len: u32) -> Result<&'a [u8]> {
        let slice = self
            .data
            .get(self.cursor..self.cursor + len as usize)
            .ok_or_else(|| anyhow!("oob"))?;
        self.cursor += len as usize;

        Ok(slice)
    }

    fn read_offsets(&mut self, width: OffsetWidth, capacity: u32) -> Result<Vec<u32>> {
        (0..capacity)
            .map(|_| self.read_offset(width))
            .collect::<Result<Vec<_>>>()
    }

    impl_read_signed_leb128!(read_i16_leb128, i16);
    impl_read_signed_leb128!(read_i32_leb128, i32);
    impl_read_signed_leb128!(read_i64_leb128, i64);
    impl_read_signed_leb128!(read_i128_leb128, i128);
}

#[macro_export]
macro_rules! impl_read_for_datatype {
    ($name:ident, $type:ty) => {
        fn $name(&mut self) -> Result<$type> {
            let width = std::mem::size_of::<$type>();
            let slice = self.read_slice(width as u32)?;
            let b = <$type>::from_le_bytes(slice.try_into()?);

            Ok(b)
        }
    };
}

#[macro_export]
macro_rules! impl_read_u24 {
    () => {
        fn read_u24(&mut self) -> Result<u32> {
            let slice = self.read_slice(3)?;
            Ok(u32::from_le_bytes([slice[0], slice[1], slice[2], 0]))
        }
    };
}

// https://github.com/rust-lang/rust/blob/30f74ff0dc4d66debc8b50724c446f817e5f75f4/compiler/rustc_serialize/src/leb128.rs
#[macro_export]
macro_rules! impl_read_signed_leb128 {
    ($name:ident, $type:ty) => {
        fn $name(&mut self) -> Result<$type> {
            let mut result = 0;
            let mut shift = 0;
            let mut byte;

            loop {
                byte = self.read_u8()?;
                result |= <$type>::from(byte & 0x7F) << shift;
                shift += 7;

                if (byte & 0x80) == 0 {
                    break;
                }
            }

            if (shift < <$type>::BITS) && ((byte & 0x40) != 0) {
                // sign extend
                result |= (!0 << shift);
            }

            Ok(result)
        }
    };
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_invalid_header() {
        let bytes = &[0b0000_0011]; // version is 3
        assert!(VariantMetadataDecoder::try_new(bytes, true).is_err());
    }

    #[test]
    fn test_simple_metadata() {
        let bytes = &[
            0b0000_0001,
            0x02,
            0x00,
            0x04,
            0x07,
            b'n',
            b'o',
            b'r',
            b'm',
            b'd',
            b'o',
            b'g',
        ];

        let mut metadata_decoder = VariantMetadataDecoder::try_new(bytes, true).unwrap();
        let metadata = metadata_decoder.read_metadata().unwrap();

        assert_eq!(metadata.header.layout, DictionaryLayout::Naive);
        assert_eq!(metadata.header.offset_width, OffsetWidth::U8);
        assert_eq!(&metadata.offsets, &[0, 4, 7]);
        assert_eq!(metadata.dictionary_len(), 2);
        assert_eq!(metadata.dictionary_strings(), &["norm", "dog"]);
    }

    #[test]
    fn test_simple_metadata_offset_width_u16() {
        let bytes = &[
            0b0100_0001,
            0x02,
            0x00,
            0x00, // offset 0
            0x00,
            0x04, // offset 4
            0x00,
            0x07, // offset 7
            0x00,
            b'n',
            b'o',
            b'r',
            b'm',
            b'd',
            b'o',
            b'g',
        ];

        let mut metadata_decoder = VariantMetadataDecoder::try_new(bytes, true).unwrap();
        let metadata = metadata_decoder.read_metadata().unwrap();

        assert_eq!(metadata.header.layout, DictionaryLayout::Naive);
        assert_eq!(metadata.header.offset_width, OffsetWidth::U16);
        assert_eq!(&metadata.offsets, &[0, 4, 7]);
        assert_eq!(metadata.dictionary_len(), 2);
        assert_eq!(metadata.dictionary_strings(), &["norm", "dog"]);
    }

    #[test]
    fn test_simple_metadata_offset_width_u24() {
        let bytes = &[
            0b1000_0001,
            0x02,
            0x00,
            0x00,
            0x00, // offset 0
            0x00,
            0x00,
            0x04, // offset 4
            0x00,
            0x00,
            0x07, // offset 7
            0x00,
            0x00,
            b'n',
            b'o',
            b'r',
            b'm',
            b'd',
            b'o',
            b'g',
        ];

        let mut metadata_decoder = VariantMetadataDecoder::try_new(bytes, true).unwrap();
        let metadata = metadata_decoder.read_metadata().unwrap();

        assert_eq!(metadata.header.layout, DictionaryLayout::Naive);
        assert_eq!(metadata.header.offset_width, OffsetWidth::U24);
        assert_eq!(&metadata.offsets, &[0, 4, 7]);
        assert_eq!(metadata.dictionary_len(), 2);
        assert_eq!(metadata.dictionary_strings(), &["norm", "dog"]);
    }

    #[test]
    fn test_metadata_invalid_offsets() {
        let bytes = &[
            0b0000_0001,
            0x02,
            0x00,
            0x07,
            0x04, // offset is not strictly increasing
            b'n',
            b'o',
            b'r',
            b'm',
            b'd',
            b'o',
            b'g',
        ];

        let mut metadata_decoder = VariantMetadataDecoder::try_new(bytes, true).unwrap();
        assert!(metadata_decoder.read_metadata().is_err());
    }

    #[test]
    fn test_metadata_with_empty_string() {
        let bytes = &[
            0b0001_0001, // sorted string
            0x02,
            0x00,
            0x00,
            0x04,
            b'n',
            b'o',
            b'r',
            b'm',
        ];

        let mut metadata_decoder = VariantMetadataDecoder::try_new(bytes, true).unwrap();
        let metadata = metadata_decoder.read_metadata().unwrap();

        assert_eq!(&metadata.offsets, &[0u32, 0u32, 4u32]);
        assert_eq!(metadata.header.layout, DictionaryLayout::SortedUnique);
        assert_eq!(metadata.header.offset_width, OffsetWidth::U8);
        assert_eq!(metadata.dictionary_len(), 2);
        assert_eq!(metadata.dictionary_strings(), &["", "norm"]);
    }

    #[test]
    fn test_metadata_with_duplicate_field_names_marked_sorted_unique() {
        let bytes = &[
            0b0001_0001, // sorted string
            0x03,
            0x00,
            0x00,
            0x00,
            0x04,
            b'n',
            b'o',
            b'r',
            b'm',
        ];

        let mut metadata_decoder = VariantMetadataDecoder::try_new(bytes, true).unwrap();
        assert!(metadata_decoder.read_metadata().is_err());
    }

    #[test]
    fn test_metadata_with_duplicate_field_names_marked_naive() {
        let bytes = &[
            0b0000_0001, // sorted string
            0x03,
            0x00,
            0x00,
            0x00,
            0x04,
            b'n',
            b'o',
            b'r',
            b'm',
        ];

        let mut metadata_decoder = VariantMetadataDecoder::try_new(bytes, true).unwrap();
        let metadata = metadata_decoder.read_metadata().unwrap();

        assert_eq!(&metadata.offsets, &[0u32, 0u32, 0, 4u32]);
        assert_eq!(metadata.header.layout, DictionaryLayout::Naive);
        assert_eq!(metadata.header.offset_width, OffsetWidth::U8);
        assert_eq!(metadata.dictionary_len(), 3);
        assert_eq!(metadata.dictionary_strings(), &["", "", "norm"]);
    }

    #[test]
    fn test_metadata_regressions() {
        let metadata_regressions = std::fs::read_dir("regressions")
            .unwrap()
            .filter_map(|entry| {
                entry
                    .map(|e| {
                        e.file_name()
                            .to_str()
                            .unwrap()
                            .starts_with("metadata")
                            .then(|| e.path())
                    })
                    .unwrap()
            })
            .collect::<Vec<_>>();

        for p in metadata_regressions {
            let data = std::fs::read(p).unwrap();
            if let Ok(mut metadata_decoder) = VariantMetadataDecoder::try_new(&data, true) {
                let res = metadata_decoder.read_metadata();
                assert!(res.is_ok() || res.is_err());
            };
        }
    }

    #[test]
    fn test_value_null() {
        let mut decoder = VariantDecoder::try_new(
            &[
                0b0000_0001,
                0x00, // no elements
                0x00, // end offset
            ],
            &[
                0, // tag 0 for null
            ],
            true,
        )
        .unwrap();

        let v = decoder.read_value().unwrap();
        assert_eq!(v, Variant::Null);
    }

    #[test]
    fn test_value_bool() {
        {
            let mut decoder = VariantDecoder::try_new(
                &[
                    0b0000_0001,
                    0x00, // no elements
                    0x00, // end offset
                ],
                &[
                    0b0000_0100, // tag 1 for null
                ],
                true,
            )
            .unwrap();
            let v = decoder.read_value().unwrap();
            assert_eq!(v, Variant::BooleanTrue);
        }

        {
            let mut decoder = VariantDecoder::try_new(
                &[
                    0b0000_0001,
                    0x00, // no elements
                    0x00, // end offset
                ],
                &[
                    0b0000_1000, // tag 0 for null
                ],
                true,
            )
            .unwrap();

            let v = decoder.read_value().unwrap();
            assert_eq!(v, Variant::BooleanFalse);
        }
    }

    #[test]
    fn test_simple_object() {
        let m = [
            0b10001, 1,  // 1 item
            0,  // start offset
            1,  // end offset
            97, // a
        ];

        let mut decoder = VariantMetadataDecoder::try_new(&m, true).unwrap();
        let metadata = decoder.read_metadata().unwrap();
        dbg!(&metadata);

        let v = [0b10, 0, 0, 1, 0];

        let mut decoder = VariantDecoder::try_new(&m, &v, true).unwrap();
        let v = decoder.read_value().unwrap();
        dbg!(&v);
    }
}