openvariant 0.1.0

use crate::{
    grammar::{OffsetWidth, Variant},
    impl_write_for_datatype, impl_write_signed_leb128,
};
use indexmap::{IndexMap, IndexSet};
use std::borrow::Cow;

#[derive(Debug, Default)]
pub struct VariantMetadataEncoder<'a> {
    field_names: IndexSet<Cow<'a, str>>,
    is_sorted: bool,
}

impl<'a> VariantMetadataEncoder<'a> {
    pub fn empty_metadata() -> Vec<u8> {
        vec![
            0b0000_0001, // version 1
            0x00,        // 0 items
            0x00,        // offset 0
        ]
    }

    pub fn is_empty(&self) -> bool {
        self.field_names.len() == 0
    }

    fn check_sorted(&mut self) {
        if self.field_names.len() == 1 {
            self.is_sorted = true;
            return;
        }

        let l = self.field_names.len();

        self.is_sorted = self.is_sorted && self.field_names[l - 2] < self.field_names[l - 1]
    }

    pub fn insert_field_name(&mut self, new_field_name: Cow<'a, str>) -> u32 {
        let (field_id, _) = self.field_names.insert_full(new_field_name);
        self.check_sorted();

        field_id as u32
    }

    pub(crate) fn find_by_field_id(&self, field_id: u32) -> &str {
        &self.field_names[field_id as usize]
    }

    pub fn finish(&self) -> Vec<u8> {
        let last_offset = self.field_names.iter().map(|f| f.len()).sum::<usize>();

        let offset_width = OffsetWidth::find_offset(last_offset as u32);

        let mut buf = ValueEncoder(Vec::with_capacity(
            1 // header byte 
            + (offset_width as usize + 1)                         // dictionary size
            + ((offset_width as usize + 1) * self.field_names.len()) // offsets
            + last_offset, // bytes
        ));

        buf.write_u8((offset_width as u8) << 6 | (self.is_sorted as u8) << 4 | 0b1);
        buf.write_offsets(offset_width, &[self.field_names.len() as u32]);

        // note this is prefix sum. this can be easily vectorized
        let mut field_offsets: Vec<u32> = Vec::with_capacity(self.field_names.len() + 1);
        field_offsets.push(0);

        for field_name in &self.field_names {
            let offset = field_name.len() as u32;
            field_offsets.push(field_offsets.last().unwrap() + offset);
        }

        buf.write_offsets(offset_width, &field_offsets);

        for field_name in &self.field_names {
            buf.write_str(field_name);
        }

        buf.finish()
    }
}

#[derive(Debug, Default)]
pub struct ValueEncoder(Vec<u8>);

impl ValueEncoder {
    pub fn write_variant<'a>(&mut self, v: &Variant<'a>) {
        match v {
            Variant::Null => self.write_primitive_header(0),
            Variant::BooleanTrue => self.write_primitive_header(1),
            Variant::BooleanFalse => self.write_primitive_header(2),
            Variant::Int8(n) => {
                self.write_primitive_header(3);
                self.write_i8(*n);
            }
            Variant::Int16(n) => {
                self.write_primitive_header(4);
                self.write_i16(*n);
            }
            Variant::Int32(n) => {
                self.write_primitive_header(5);
                self.write_i32(*n);
            }
            Variant::Int64(n) => {
                self.write_primitive_header(6);
                self.write_i64(*n);
            }
            Variant::Double(f) => {
                self.write_primitive_header(7);
                self.write_f64(*f);
            }
            Variant::Decimal4(precision, scale) => {
                self.write_primitive_header(8);
                self.write_u8(*precision);
                self.write_i32(*scale);
            }
            Variant::Decimal8(precision, scale) => {
                self.write_primitive_header(9);
                self.write_u8(*precision);
                self.write_i64(*scale);
            }
            Variant::Decimal16(precision, scale) => {
                self.write_primitive_header(10);
                self.write_u8(*precision);
                self.write_i128(*scale);
            }
            Variant::Date(d) => {
                self.write_primitive_header(11);
                self.write_u32(*d);
            }
            Variant::TimestampMicros(t) => {
                self.write_primitive_header(12);
                self.write_u64(*t);
            }
            Variant::TimestampNTZMicros(t) => {
                self.write_primitive_header(13);
                self.write_u64(*t);
            }
            Variant::Float(f) => {
                self.write_primitive_header(14);
                self.write_f32(*f);
            }
            Variant::Binary(cow) => {
                self.write_primitive_header(15);
                self.0.extend_from_slice(cow);
            }
            Variant::String(cow) => {
                self.write_primitive_header(16);
                self.write_str(cow);
            }
            Variant::TimeNTZMicros(t) => {
                self.write_primitive_header(17);
                self.write_u64(*t);
            }
            Variant::TimestampNanos(t) => {
                self.write_primitive_header(18);
                self.write_u64(*t);
            }
            Variant::TimestampNTZNanos(t) => {
                self.write_primitive_header(19);
                self.write_u64(*t);
            }
            Variant::UUID(u) => {
                self.write_primitive_header(20);
                self.write_u128(*u);
            }
            Variant::Int16Leb128(i) => {
                self.write_primitive_header(21);
                self.write_i16_leb128(*i);
            }
            Variant::Int32Leb128(i) => {
                self.write_primitive_header(22);
                self.write_i32_leb128(*i);
            }
            Variant::Int64Leb128(i) => {
                self.write_primitive_header(23);
                self.write_i64_leb128(*i);
            }
            Variant::Int128Leb128(i) => {
                self.write_primitive_header(24);
                self.write_i128_leb128(*i);
            }
            Variant::ShortString(short_string) => {
                self.write_short_string_header(short_string.len());
                self.write_str(short_string.inner());
            }
            Variant::Object(_variant_object) => {
                // self.write_object_header(&variant_object);
                todo!();
            }
            Variant::Array(_array) => {
                // self.write_array_header(&array);
                todo!();
            }
        }
    }

    pub fn finish(self) -> Vec<u8> {
        self.0
    }

    pub fn append(&mut self, other: Self) {
        self.0.extend_from_slice(&other.0);
    }

    fn write_primitive_header(&mut self, tag: u8) {
        self.0.push(tag << 2);
    }

    fn write_short_string_header(&mut self, len: u8) {
        self.0.push(len << 2 | 1)
    }

    fn write_object_header(
        &mut self,
        num_fields: usize,
        largest_field_id: OffsetWidth,
        largest_field_offset: OffsetWidth,
    ) {
        let is_large = num_fields > u8::MAX as usize;
        let field_id_size_minus_one = largest_field_id as u8;
        let field_offset_size_minus_one = largest_field_offset as u8;

        let value_header =
            (is_large as u8) << 4 | field_id_size_minus_one << 2 | field_offset_size_minus_one;

        self.0.push(value_header << 2 | 2);

        if is_large {
            self.write_u32(num_fields as u32);
        } else {
            self.0.push(num_fields as u8)
        };
    }

    fn write_array_header(&mut self, len: u32, largest_field_offset: OffsetWidth) {
        let is_large = len > u8::MAX as u32;
        let field_offset_size_minus_one = largest_field_offset as u8;

        let value_header = (is_large as u8) << 2 | field_offset_size_minus_one;

        self.0.push(value_header << 2 | 3);
        if is_large {
            self.write_u32(len);
        } else {
            self.0.push(len as u8);
        }
    }

    fn write_offsets<'a, I>(&mut self, offset_width: OffsetWidth, offsets: I)
    where
        I: IntoIterator<Item = &'a u32>,
    {
        match offset_width {
            OffsetWidth::U8 => self.0.extend(offsets.into_iter().map(|&o| o as u8)),
            OffsetWidth::U16 => self
                .0
                .extend(offsets.into_iter().flat_map(|&o| (o as u16).to_le_bytes())),
            OffsetWidth::U24 => self.0.extend(offsets.into_iter().flat_map(|&o| {
                let mut temp_buf = [0u8; 3];
                temp_buf.copy_from_slice(&o.to_le_bytes()[..3]);
                temp_buf
            })),
            OffsetWidth::U32 => self
                .0
                .extend(offsets.into_iter().flat_map(|&o| o.to_le_bytes())),
        }
    }

    fn write_str(&mut self, str: &str) {
        self.0.extend_from_slice(str.as_bytes());
    }

    impl_write_for_datatype!(write_i8, i8);
    impl_write_for_datatype!(write_i16, i16);
    impl_write_for_datatype!(write_i32, i32);
    impl_write_for_datatype!(write_i64, i64);
    impl_write_for_datatype!(write_i128, i128);

    impl_write_for_datatype!(write_f32, f32);
    impl_write_for_datatype!(write_f64, f64);

    impl_write_for_datatype!(write_u8, u8);
    impl_write_for_datatype!(write_u32, u32);
    impl_write_for_datatype!(write_u64, u64);
    impl_write_for_datatype!(write_u128, u128);

    impl_write_signed_leb128!(write_i16_leb128, i16);
    impl_write_signed_leb128!(write_i32_leb128, i32);
    impl_write_signed_leb128!(write_i64_leb128, i64);
    impl_write_signed_leb128!(write_i128_leb128, i128);
}

#[derive(Debug)]
pub struct VariantEncoder<'a, 'b> {
    metadata_encoder: &'a mut VariantMetadataEncoder<'b>,
    value_encoder: ValueEncoder,
}

impl<'a, 'b> VariantEncoder<'a, 'b> {
    pub fn new(metadata_encoder: &'a mut VariantMetadataEncoder<'b>) -> Self {
        Self {
            metadata_encoder,
            value_encoder: Default::default(),
        }
    }

    pub fn write_variant<T>(&mut self, v: T)
    where
        T: Into<Variant<'b>>,
    {
        let v = v.into();

        self.value_encoder.write_variant(&v);

        // for complex types we should update the metadata
        // match v {
        //     Variant::Array(variant_array) => {}
        //     Variant::Object(variant_object) => {}
        //     _ => {}
        // }
    }

    pub fn write_object(&mut self) -> ObjectEncoder<'_, 'b> {
        ObjectEncoder::new(&mut self.metadata_encoder, &mut self.value_encoder)
    }

    pub fn write_array(&mut self) -> ArrayEncoder<'_, 'b> {
        ArrayEncoder::new(&mut self.metadata_encoder, &mut self.value_encoder)
    }

    pub fn finish(self) -> (Vec<u8>, Vec<u8>) {
        (self.metadata_encoder.finish(), self.value_encoder.finish())
    }
}

#[derive(Debug)]
pub struct ArrayEncoder<'a, 'b> {
    metadata: &'a mut VariantMetadataEncoder<'b>,
    parent_value_encoder: &'a mut ValueEncoder,

    offsets: Vec<u32>,
    value_encoder: ValueEncoder,
}

impl<'a, 'b> ArrayEncoder<'a, 'b> {
    pub fn new(
        metadata: &'a mut VariantMetadataEncoder<'b>,
        parent_value_encoder: &'a mut ValueEncoder,
    ) -> Self {
        Self {
            metadata,
            parent_value_encoder,
            offsets: Default::default(),
            value_encoder: ValueEncoder::default(),
        }
    }

    pub fn append<T>(&mut self, variant: T)
    where
        T: Into<Variant<'b>>,
    {
        self.offsets.push(self.value_encoder.0.len() as u32);
        self.value_encoder.write_variant(&variant.into());
    }

    pub fn insert_object<T>(&'b mut self) -> ObjectEncoder<'b, 'b> {
        ObjectEncoder::new(self.metadata, &mut self.value_encoder)
    }

    pub fn insert_array<T>(&'b mut self) -> ArrayEncoder<'b, 'b> {
        ArrayEncoder::new(self.metadata, &mut self.value_encoder)
    }

    pub fn finish(self) {
        let largest_field_offset_width =
            OffsetWidth::find_offset(self.value_encoder.0.len() as u32);

        self.parent_value_encoder
            .write_array_header(self.offsets.len() as u32, largest_field_offset_width);

        self.parent_value_encoder
            .write_offsets(largest_field_offset_width, &self.offsets);
        self.parent_value_encoder.write_offsets(
            largest_field_offset_width,
            &[self.value_encoder.0.len() as u32],
        );

        self.parent_value_encoder.append(self.value_encoder);
    }
}

#[derive(Debug)]
pub struct ObjectEncoder<'a, 'b> {
    metadata: &'a mut VariantMetadataEncoder<'b>,
    parent_value_encoder: &'a mut ValueEncoder,

    fields: IndexMap<u32, u32>, // field id to field offset
    largest_field_id: u32,

    value_encoder: ValueEncoder,
}

impl<'a, 'b> ObjectEncoder<'a, 'b> {
    pub fn new(
        metadata: &'a mut VariantMetadataEncoder<'b>,
        parent_value_encoder: &'a mut ValueEncoder,
    ) -> Self {
        Self {
            metadata,
            parent_value_encoder,
            fields: Default::default(),
            largest_field_id: Default::default(),
            value_encoder: ValueEncoder::default(),
        }
    }

    fn add_field_name(&mut self, field_name: Cow<'b, str>) {
        // update the metadata
        let field_id = self.metadata.insert_field_name(field_name);

        // this is a design choice, if a duplicate key gets inserted, we write over it
        self.fields
            .insert(field_id, self.value_encoder.0.len() as u32);

        self.largest_field_id = self.largest_field_id.max(field_id);
    }

    pub fn insert<K, V>(&mut self, field_name: K, variant: V)
    where
        K: Into<Cow<'b, str>>,
        V: Into<Variant<'b>>,
    {
        self.add_field_name(field_name.into());
        self.value_encoder.write_variant(&variant.into());
    }

    pub fn insert_object<K, V>(&'b mut self, field_name: K) -> ObjectEncoder<'b, 'b>
    where
        K: Into<Cow<'b, str>>,
    {
        self.add_field_name(field_name.into());
        ObjectEncoder::new(self.metadata, &mut self.value_encoder)
    }

    pub fn insert_array<K>(&'b mut self, field_name: K) -> ArrayEncoder<'b, 'b>
    where
        K: Into<Cow<'b, str>>,
    {
        self.add_field_name(field_name.into());
        ArrayEncoder::new(self.metadata, &mut self.value_encoder)
    }

    pub fn finish(self) {
        // write out the header
        let largest_field_id_width = OffsetWidth::find_offset(self.largest_field_id);
        let largest_field_offset_width =
            OffsetWidth::find_offset(self.value_encoder.0.len() as u32);

        self.parent_value_encoder.write_object_header(
            self.fields.len(),
            largest_field_id_width,
            largest_field_offset_width,
        );

        let mut fields = self.fields;

        if self.metadata.is_sorted {
            // at the time when we write out this object,  if the metadata indicates it's sorted
            // we just need to order by field id
            fields.sort_by(|field_id1, _, field_id2, _| field_id1.cmp(field_id2));
        } else {
            // we need to order by the respective field name
            fields.sort_by(|field_id1, _, field_id2, _| {
                self.metadata
                    .find_by_field_id(*field_id1)
                    .cmp(self.metadata.find_by_field_id(*field_id2))
            });
        }

        let field_ids = fields.keys();
        self.parent_value_encoder
            .write_offsets(largest_field_id_width, field_ids);

        let field_offsets = fields.values();
        self.parent_value_encoder
            .write_offsets(largest_field_offset_width, field_offsets);
        self.parent_value_encoder.write_offsets(
            largest_field_offset_width,
            &[self.value_encoder.0.len() as u32],
        );

        self.parent_value_encoder.append(self.value_encoder);
    }
}

#[macro_export]
macro_rules! impl_write_for_datatype {
    ($name:ident, $type:ty) => {
        fn $name(&mut self, n: $type) {
            self.0.extend(n.to_le_bytes());
        }
    };
}

// https://github.com/rust-lang/rust/blob/30f74ff0dc4d66debc8b50724c446f817e5f75f4/compiler/rustc_serialize/src/leb128.rs
#[macro_export]
macro_rules! impl_write_signed_leb128 {
    ($name:ident, $type:ty) => {
        fn $name(&mut self, mut value: $type) {
            let mut out = [0u8; (std::mem::size_of::<$type>() * 8).div_ceil(7)];
            let mut i = 0;

            loop {
                let mut byte = (value as u8) & 0x7f;
                value >>= 7;
                let more = !(((value == 0) && ((byte & 0x40) == 0))
                    || ((value == -1) && ((byte & 0x40) != 0)));

                if more {
                    byte |= 0x80; // Mark this byte to show that more bytes will follow.
                }

                out[i] = byte;
                i += 1;

                if !more {
                    break;
                }
            }

            self.0.extend(&out[..i]);
        }
    };
}

#[cfg(test)]
mod tests {
    use crate::{
        VariantDecoder, VariantMetadataDecoder,
        grammar::{DictionaryLayout, VariantArray, VariantObject},
    };

    use super::*;

    #[test]
    fn test_build_basic_null() {
        let mut meta_encoder = VariantMetadataEncoder::default();
        let mut encoder = VariantEncoder::new(&mut meta_encoder);
        encoder.write_variant(());
        let (m, v) = encoder.finish();

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

        assert_eq!(v, Variant::Null);
    }

    #[test]
    fn test_build_basic_bool() {
        let mut meta_encoder = VariantMetadataEncoder::default();
        let mut encoder = VariantEncoder::new(&mut meta_encoder);

        encoder.write_variant(false);
        let (m, v) = encoder.finish();

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

        assert_eq!(v, Variant::BooleanFalse);
    }

    #[test]
    fn test_build_basic_int8() {
        let mut meta_encoder = VariantMetadataEncoder::default();
        let mut encoder = VariantEncoder::new(&mut meta_encoder);

        encoder.write_variant(32i8);
        let (m, v) = encoder.finish();

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

        assert_eq!(v, Variant::Int8(32));
    }

    #[test]
    fn test_build_basic_object() {
        let mut meta_encoder = VariantMetadataEncoder::default();
        let mut encoder = VariantEncoder::new(&mut meta_encoder);

        {
            let mut obj_encoder = encoder.write_object();
            obj_encoder.insert("a", ());
            obj_encoder.finish();
        }
        let (m, v) = encoder.finish();

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

        assert_eq!(
            v,
            Variant::Object(VariantObject {
                field_ids: vec![0],
                field_offsets: vec![0, 1],
                values: vec![Variant::Null]
            })
        );
    }

    #[test]
    fn test_build_object_reverse_field_order() -> anyhow::Result<()> {
        let mut meta_encoder = VariantMetadataEncoder::default();
        let mut encoder = VariantEncoder::new(&mut meta_encoder);

        {
            let mut obj_encoder = encoder.write_object();
            obj_encoder.insert("c", Variant::Int32Leb128(4));
            obj_encoder.insert("b", Variant::Int32Leb128(127));
            obj_encoder.insert("a", Variant::Int32Leb128(255));
            obj_encoder.finish();
        }

        assert!(!encoder.metadata_encoder.is_sorted);

        let (m, v) = encoder.finish();

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

        assert_eq!(metadata.dictionary_strings(), &["c", "b", "a"]);
        assert_eq!(metadata.header.layout, DictionaryLayout::Naive);

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

        assert_eq!(
            v,
            Variant::Object(VariantObject {
                field_ids: vec![2, 1, 0],
                field_offsets: vec![5, 2, 0, 8],
                values: vec![
                    Variant::Int32Leb128(255),
                    Variant::Int32Leb128(127),
                    Variant::Int32Leb128(4),
                ]
            })
        );

        Ok(())
    }

    #[test]
    fn test_build_leb128() {
        let mut meta_encoder = VariantMetadataEncoder::default();
        let mut encoder = VariantEncoder::new(&mut meta_encoder);

        encoder.write_variant(Variant::Int32Leb128(128));
        let (m, v) = encoder.finish();

        /*
        A value buffer with a Int32 would have looked like this:
        &v = [
            20,
            128,
            0,
            0,
            0,
        ]

        meanwhile with leb128:
        &v = [
            88,
            128,
            1
        ]
        */

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

        assert_eq!(v, Variant::Int32Leb128(128));
    }

    #[test]
    fn test_build_basic_array() {
        let mut meta_encoder = VariantMetadataEncoder::default();
        let mut encoder = VariantEncoder::new(&mut meta_encoder);

        {
            let mut arr_encoder = encoder.write_array();

            arr_encoder.append(Variant::Null);
            arr_encoder.append(Variant::Int16Leb128(20));
            arr_encoder.append(Variant::Int32Leb128(127));
            arr_encoder.append(Variant::Int64Leb128(1));

            arr_encoder.finish();
        }

        let (m, v) = encoder.finish();

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

        assert_eq!(
            v,
            Variant::Array(VariantArray {
                field_offsets: vec![0, 1, 3, 6, 8],
                values: vec![
                    Variant::Null,
                    Variant::Int16Leb128(20),
                    Variant::Int32Leb128(127),
                    Variant::Int64Leb128(1)
                ],
            })
        )
    }

    #[test]
    fn test_nested_object() {}
}