fory-core 1.6.0-rc.1

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

use std::collections::BTreeMap;
use std::marker::PhantomData;

use crate::buffer::Writer;
use crate::error::Error;

use super::bit_util::{bitmap_width, round_up_to_word, set_bit, slot_width};
use super::row::{Row, RowValue};

enum WriteTarget<'a, 'b> {
    Fixed(&'a mut [u8]),
    Variable(&'a mut Writer<'b>),
}

/// The exact destination for one Row Format value.
///
/// This type is public only because `ForyRow` implementations are generated in
/// downstream crates. Applications should use [`to_row`] or [`to_row_into`]
/// instead.
#[doc(hidden)]
pub struct ValueWriter<'a, 'b> {
    target: WriteTarget<'a, 'b>,
}

impl<'a, 'b> ValueWriter<'a, 'b> {
    pub(crate) fn fixed(bytes: &'a mut [u8]) -> Self {
        Self {
            target: WriteTarget::Fixed(bytes),
        }
    }

    pub(crate) fn variable(writer: &'a mut Writer<'b>) -> Self {
        Self {
            target: WriteTarget::Variable(writer),
        }
    }

    /// Writes the complete bytes of a scalar, string, or binary value.
    pub fn write_bytes(self, bytes: &[u8]) -> Result<(), Error> {
        match self.target {
            WriteTarget::Fixed(output) => {
                if output.len() != bytes.len() {
                    return Err(Error::invalid_data("row fixed-width value size mismatch"));
                }
                output.copy_from_slice(bytes);
            }
            WriteTarget::Variable(writer) => {
                writer.write_bytes(bytes);
            }
        }
        Ok(())
    }

    /// Starts a derived struct at this value's variable destination.
    pub fn struct_writer(self, num_fields: usize) -> Result<StructWriter<'a, 'b>, Error> {
        StructWriter::new(num_fields, self.into_variable()?)
    }

    pub(crate) fn into_variable(self) -> Result<&'a mut Writer<'b>, Error> {
        match self.target {
            WriteTarget::Variable(writer) => Ok(writer),
            WriteTarget::Fixed(_) => Err(Error::invalid_data(
                "variable row value cannot use a fixed-width slot",
            )),
        }
    }
}

/// Writes one standard Row Format struct into a shared root buffer.
///
/// This type is public only for code generated by `ForyRow`.
#[doc(hidden)]
pub struct StructWriter<'a, 'b> {
    writer: &'a mut Writer<'b>,
    base_offset: usize,
    bitmap_width: usize,
    num_fields: usize,
    fixed_end: usize,
}

impl<'a, 'b> StructWriter<'a, 'b> {
    fn new(num_fields: usize, writer: &'a mut Writer<'b>) -> Result<Self, Error> {
        let base_offset = writer.len();
        let bitmap_width = bitmap_width(num_fields)?;
        let slots_size = num_fields
            .checked_mul(8)
            .ok_or_else(|| Error::invalid_data("row fixed region size overflow"))?;
        let fixed_size = bitmap_width
            .checked_add(slots_size)
            .ok_or_else(|| Error::invalid_data("row fixed region size overflow"))?;
        let fixed_end = base_offset
            .checked_add(fixed_size)
            .ok_or_else(|| Error::invalid_data("row fixed region size overflow"))?;
        writer.skip(fixed_size);
        Ok(Self {
            writer,
            base_offset,
            bitmap_width,
            num_fields,
            fixed_end,
        })
    }

    /// Writes a field at its schema ordinal.
    pub fn write<T: RowValue + ?Sized>(&mut self, index: usize, value: &T) -> Result<(), Error> {
        if index >= self.num_fields {
            return Err(Error::buffer_out_of_bound(index, 1, self.num_fields));
        }
        if value.is_null() {
            let bitmap =
                &mut self.writer.bf[self.base_offset..self.base_offset + self.bitmap_width];
            set_bit(bitmap, index);
            return Ok(());
        }

        let slot_offset = self
            .base_offset
            .checked_add(self.bitmap_width)
            .and_then(|offset| offset.checked_add(index * 8))
            .ok_or_else(|| Error::invalid_data("row field offset overflow"))?;
        match T::FIXED_SIZE {
            Some(width) => {
                slot_width(Some(width))?;
                let capacity = self.writer.bf.len();
                let output = self
                    .writer
                    .bf
                    .get_mut(slot_offset..slot_offset + width)
                    .ok_or_else(|| Error::buffer_out_of_bound(slot_offset, width, capacity))?;
                value.write(ValueWriter::fixed(output))
            }
            None => write_variable(
                self.writer,
                self.base_offset,
                self.fixed_end,
                slot_offset,
                value,
            ),
        }
    }
}

pub(crate) struct ArrayWriter<'a, 'b, T: RowValue> {
    writer: &'a mut Writer<'b>,
    base_offset: usize,
    bitmap_width: usize,
    header_size: usize,
    num_elements: usize,
    element_size: usize,
    fixed_end: usize,
    marker: PhantomData<T>,
}

impl<'a, 'b, T: RowValue> ArrayWriter<'a, 'b, T> {
    pub(crate) fn new(num_elements: usize, writer: &'a mut Writer<'b>) -> Result<Self, Error> {
        let count = u64::try_from(num_elements)
            .map_err(|_| Error::invalid_data("row array element count exceeds u64"))?;
        let base_offset = writer.len();
        let bitmap_width = bitmap_width(num_elements)?;
        let header_size = 8usize
            .checked_add(bitmap_width)
            .ok_or_else(|| Error::invalid_data("row array header size overflow"))?;
        let element_size = slot_width(T::FIXED_SIZE)?;
        let element_bytes = num_elements
            .checked_mul(element_size)
            .ok_or_else(|| Error::invalid_data("row array fixed region size overflow"))?;
        let aligned_element_bytes = round_up_to_word(element_bytes)?;
        let fixed_size = header_size
            .checked_add(aligned_element_bytes)
            .ok_or_else(|| Error::invalid_data("row array fixed region size overflow"))?;
        let fixed_end = base_offset
            .checked_add(fixed_size)
            .ok_or_else(|| Error::invalid_data("row array fixed region size overflow"))?;

        writer.write_u64(count);
        writer.skip(fixed_size - 8);
        Ok(Self {
            writer,
            base_offset,
            bitmap_width,
            header_size,
            num_elements,
            element_size,
            fixed_end,
            marker: PhantomData,
        })
    }

    pub(crate) fn write(&mut self, index: usize, value: &T) -> Result<(), Error> {
        if index >= self.num_elements {
            return Err(Error::buffer_out_of_bound(index, 1, self.num_elements));
        }
        if value.is_null() {
            let bitmap_start = self.base_offset + 8;
            let bitmap = &mut self.writer.bf[bitmap_start..bitmap_start + self.bitmap_width];
            set_bit(bitmap, index);
            return Ok(());
        }

        let slot_offset = self
            .base_offset
            .checked_add(self.header_size)
            .and_then(|offset| offset.checked_add(index * self.element_size))
            .ok_or_else(|| Error::invalid_data("row array element offset overflow"))?;
        match T::FIXED_SIZE {
            Some(width) => {
                let capacity = self.writer.bf.len();
                let output = self
                    .writer
                    .bf
                    .get_mut(slot_offset..slot_offset + width)
                    .ok_or_else(|| Error::buffer_out_of_bound(slot_offset, width, capacity))?;
                value.write(ValueWriter::fixed(output))
            }
            None => write_variable(
                self.writer,
                self.base_offset,
                self.fixed_end,
                slot_offset,
                value,
            ),
        }
    }
}

pub(crate) struct MapWriter<'a, 'b> {
    writer: &'a mut Writer<'b>,
    base_offset: usize,
}

impl<'a, 'b> MapWriter<'a, 'b> {
    pub(crate) fn new(writer: &'a mut Writer<'b>) -> Self {
        let base_offset = writer.len();
        writer.skip(8);
        Self {
            writer,
            base_offset,
        }
    }

    pub(crate) fn write<K, V>(&mut self, values: &BTreeMap<K, V>) -> Result<(), Error>
    where
        K: RowValue + Ord,
        V: RowValue,
    {
        let keys_start = self.writer.len();
        {
            let mut keys = ArrayWriter::<K>::new(values.len(), self.writer)?;
            for (index, key) in values.keys().enumerate() {
                keys.write(index, key)?;
            }
        }
        let keys_size = self
            .writer
            .len()
            .checked_sub(keys_start)
            .ok_or_else(|| Error::invalid_data("row map key array size underflow"))?;
        let keys_size = u64::try_from(keys_size)
            .map_err(|_| Error::invalid_data("row map key array exceeds u64"))?;
        self.writer
            .set_bytes(self.base_offset, &keys_size.to_le_bytes());

        let mut items = ArrayWriter::<V>::new(values.len(), self.writer)?;
        for (index, value) in values.values().enumerate() {
            items.write(index, value)?;
        }
        Ok(())
    }
}

fn write_variable<T: RowValue + ?Sized>(
    writer: &mut Writer<'_>,
    base_offset: usize,
    fixed_end: usize,
    slot_offset: usize,
    value: &T,
) -> Result<(), Error> {
    let data_start = writer.len();
    if data_start < fixed_end {
        return Err(Error::invalid_data(
            "row variable data overlaps the fixed region",
        ));
    }
    let relative_offset = data_start
        .checked_sub(base_offset)
        .ok_or_else(|| Error::invalid_data("row variable offset underflow"))?;
    let relative_offset = u32::try_from(relative_offset)
        .map_err(|_| Error::invalid_data("row variable offset exceeds u32"))?;

    value.write(ValueWriter::variable(writer))?;
    let size = writer
        .len()
        .checked_sub(data_start)
        .ok_or_else(|| Error::invalid_data("row variable size underflow"))?;
    let wire_size =
        u32::try_from(size).map_err(|_| Error::invalid_data("row variable size exceeds u32"))?;
    let aligned_size = round_up_to_word(size)?;
    writer.skip(aligned_size - size);

    let offset_and_size = ((relative_offset as u64) << 32) | wire_size as u64;
    writer.set_bytes(slot_offset, &offset_and_size.to_le_bytes());
    Ok(())
}

/// Encodes a struct, array, or map as a Standard Row Format root.
pub fn to_row<T: Row + ?Sized>(value: &T) -> Result<Vec<u8>, Error> {
    let mut buffer = Vec::new();
    to_row_into(value, &mut buffer)?;
    Ok(buffer)
}

/// Replaces `buffer` with one encoded Standard Row Format root.
///
/// Existing capacity is retained for reuse. If encoding fails, `buffer` is
/// left empty.
pub fn to_row_into<T: Row + ?Sized>(value: &T, buffer: &mut Vec<u8>) -> Result<(), Error> {
    buffer.clear();
    let result = {
        let mut writer = Writer::from_buffer(buffer);
        value.write(ValueWriter::variable(&mut writer))
    };
    if result.is_err() {
        buffer.clear();
    }
    result
}