arcium-primitives 0.8.0

//! Fast, architecture-independent fixed-size (de)serialization.
//!
//! [`InPlaceCodec`] encodes a value to (and decodes it from) a fixed-size byte window directly,
//! bypassing the per-element `serde` dispatch. The encoding is deliberately *not*
//! required to match any external format (`serde`/`bincode`); it only has to be
//! architecture-independent (fixed-width little-endian) and fast.

use std::mem::MaybeUninit;

use crate::errors::PrimitiveError;

pub mod bincode_io;
pub mod containers;
pub mod prims;

/// A type that can be (de)serialized to/from a fixed-size byte window directly.
///
/// The encoding must be *architecture-independent* (same bytes on any target — in practice
/// fixed-width little-endian) but is otherwise this codec's own format: it need **not** match the
/// type's `serde::Serialize`/`Deserialize` nor its `bincode` encoding.
///
/// # Safety
/// Implementors must guarantee:
/// - `write_le_bytes` initializes every one of the `ENCODED_SIZE` bytes of `out`, without reading
///   `out` (it may be uninitialized memory).
/// - `read_le_bytes` is unbiased: given `bytes.len() == ENCODED_SIZE`, if `write_le_bytes` could
///   have produced `bytes`, `read_le_bytes` must return the same value.
pub unsafe trait InPlaceCodec: Sized {
    /// Encoded width in bytes.
    const ENCODED_SIZE: usize;

    /// Write `self`'s canonical encoding into `out`, initializing every byte.
    /// `out.len() == Self::ENCODED_SIZE`.
    fn write_le_bytes(&self, out: &mut [MaybeUninit<u8>]);

    /// Parse `self` from `bytes`, validating it. `bytes.len() == Self::ENCODED_SIZE`.
    /// Returns an error on an invalid (e.g. non-canonical) encoding.
    fn read_le_bytes(bytes: &[u8]) -> Result<Self, PrimitiveError>;

    /// Number of consecutive elements that encode together as a "pack". `1` (the default) means no
    /// grouping — each element is encoded independently via [`Self::write_le_bytes`]. Must be `>=
    /// 1`.
    ///
    /// Container codecs (e.g. `HeapArray<_, N>`) split a run of elements into `N / PACK` full packs
    /// plus an `N % PACK` per-element tail. Types with `PACK > 1` must be `Copy` (see
    /// [`Self::read_pack`]).
    const PACK: usize = 1;

    /// Bytes a full `PACK`-element pack occupies. Defaults to `PACK * ENCODED_SIZE` (the packed
    /// bytes are exactly the per-element bytes laid out back-to-back — the case for a pure
    /// vectorization like `Mersenne107`). Override with a smaller value for a sub-byte packing
    /// (e.g. `Gf2`: `PACK = 8`, `PACK_BYTES = 1`).
    const PACK_BYTES: usize = Self::PACK * Self::ENCODED_SIZE;

    /// Encode exactly [`Self::PACK`] elements (`items.len() == PACK`) into `out`
    /// (`out.len() == PACK_BYTES`), initializing every byte.
    ///
    /// The default encodes the group element-by-element (valid whenever
    /// `PACK_BYTES == PACK * ENCODED_SIZE`); override it for a vectorized or sub-byte-packed
    /// encoding.
    fn write_pack(items: &[Self], out: &mut [MaybeUninit<u8>]) {
        for (chunk, item) in out.chunks_exact_mut(Self::ENCODED_SIZE).zip(items) {
            item.write_le_bytes(chunk);
        }
    }

    /// Decode a full pack: read [`Self::PACK`] elements from `bytes` (`bytes.len() == PACK_BYTES`)
    /// into `out` (`out.len() == PACK`), returning an error on an invalid encoding.
    ///
    /// The default decodes element-by-element (valid whenever `PACK_BYTES == PACK * ENCODED_SIZE`).
    /// On an error it may leave some of `out` written; the container caller drops only *completed*
    /// packs, so any type with `PACK > 1` must be `Copy` (partially-written packs are not dropped).
    fn read_pack(bytes: &[u8], out: &mut [MaybeUninit<Self>]) -> Result<(), PrimitiveError> {
        for (chunk, slot) in bytes.chunks_exact(Self::ENCODED_SIZE).zip(out.iter_mut()) {
            slot.write(Self::read_le_bytes(chunk)?);
        }
        Ok(())
    }

    /// Serialize `self` into a freshly-allocated buffer of exactly `Self::ENCODED_SIZE` bytes,
    /// bypassing per-element `serde` dispatch.
    fn to_inplace_bytes(&self) -> Vec<u8> {
        let mut out = Vec::with_capacity(Self::ENCODED_SIZE);
        let spare = &mut out.spare_capacity_mut()[..Self::ENCODED_SIZE];
        self.write_le_bytes(spare);
        // SAFETY: `write_le_bytes` initializes every byte of `spare` above.
        unsafe { out.set_len(Self::ENCODED_SIZE) };
        out
    }

    /// Deserialize `Self` from exactly `Self::ENCODED_SIZE` bytes.
    fn from_inplace_bytes(bytes: &[u8]) -> Result<Self, PrimitiveError> {
        if bytes.len() != Self::ENCODED_SIZE {
            return Err(PrimitiveError::InvalidSize(Self::ENCODED_SIZE, bytes.len()));
        }
        Self::read_le_bytes(bytes)
    }
}

/// Whether a trailing partial group of `rem` (`0 < rem < PACK`) elements is smaller encoded
/// individually (`rem * ENCODED_SIZE` bytes) than padded up to a full pack (`PACK_BYTES` bytes).
/// `rem` is a plain count, never the elements' values, so the encoder and decoder always agree on
/// this without any extra tag on the wire — see [`write_packed_le_bytes`]/[`read_packed_le_bytes`].
const fn tail_is_unpacked<T: InPlaceCodec>(rem: usize) -> bool {
    rem * T::ENCODED_SIZE < T::PACK_BYTES
}

/// Bytes needed for a trailing partial group of `rem` (`0 <= rem < PACK`) elements: whichever of
/// padding to a full pack or encoding `rem` elements individually is smaller (see
/// [`tail_is_unpacked`]).
const fn tail_size<T: InPlaceCodec>(rem: usize) -> usize {
    if rem == 0 {
        0
    } else if tail_is_unpacked::<T>(rem) {
        rem * T::ENCODED_SIZE
    } else {
        T::PACK_BYTES
    }
}

/// Shared by `HeapArray<T, M>`, `hybrid_array::Array<T, N>` and `[T; N]`'s `InPlaceCodec` impls:
/// the `PACK`-aware encoded size, `(len / PACK) * PACK_BYTES + tail_size(len % PACK)` — see
/// [`tail_size`]. For `PACK == 1` this reduces to `T::ENCODED_SIZE * len`.
pub(crate) const fn packed_size<T: InPlaceCodec>(len: usize) -> usize {
    (len / T::PACK) * T::PACK_BYTES + tail_size::<T>(len % T::PACK)
}

/// Shared by `HeapArray<T, M>`, `hybrid_array::Array<T, N>` and `[T; N]`'s `InPlaceCodec` impls:
/// write `items` as `items.len() / PACK` full packs (each `PACK_BYTES`), plus a trailing
/// `items.len() % PACK`-element tail encoded however [`tail_size`] says is smaller — padded up to
/// one more full pack, or element-by-element. `out.len()` must equal
/// `packed_size::<T>(items.len())`.
///
/// `T::PACK` is a constant, so the `PACK == 1` branch is compile-time-folded to a single flat loop
/// with no per-element `write_pack` call / iterator setup — the common case (field elements,
/// points, ...).
pub(crate) fn write_packed_le_bytes<T: InPlaceCodec>(items: &[T], out: &mut [MaybeUninit<u8>]) {
    if T::PACK == 1 {
        for (chunk, elem) in out.chunks_exact_mut(T::ENCODED_SIZE).zip(items) {
            elem.write_le_bytes(chunk);
        }
        return;
    }
    let n_full = items.len() / T::PACK;
    let rem = items.len() % T::PACK;
    let (packed, last) = out.split_at_mut(n_full * T::PACK_BYTES);
    for (chunk, group) in packed
        .chunks_exact_mut(T::PACK_BYTES)
        .zip(items.chunks_exact(T::PACK))
    {
        T::write_pack(group, chunk);
    }
    if rem == 0 {
        return;
    }
    let tail_items = &items[n_full * T::PACK..];
    if tail_is_unpacked::<T>(rem) {
        // Encoding the tail element-by-element is smaller than padding it to a full pack.
        for (chunk, item) in last.chunks_exact_mut(T::ENCODED_SIZE).zip(tail_items) {
            item.write_le_bytes(chunk);
        }
        return;
    }
    // A trailing partial group is padded with duplicates of its last element up to a full
    // `PACK`-element pack, so `write_pack` — whose bit-packed layout depends on every
    // element's value, not just the count — can encode it directly. The dummy values are
    // never read back: `read_packed_le_bytes` decodes the padded pack and keeps only its
    // first `rem` elements.
    //
    // SAFETY: `T::PACK > 1` requires `T: Copy` (see `read_pack`'s doc), so duplicating an
    // element's bytes via `ptr::read` and letting the duplicate drop normally alongside the
    // original (still owned by `items`) is exactly `Copy` semantics — no double-free.
    let mut padded: Vec<T> = tail_items
        .iter()
        .map(|item| unsafe { std::ptr::read(item) })
        .collect();
    padded.resize_with(T::PACK, || unsafe { std::ptr::read(&tail_items[rem - 1]) });
    T::write_pack(&padded, last);
}

/// Shared by `HeapArray<T, M>` and `hybrid_array::Array<T, N>`'s `InPlaceCodec` impls: the
/// decoding counterpart of [`write_packed_le_bytes`]. Fills every one of `data`'s `data.len()`
/// slots from `bytes` (`bytes.len()` must equal `packed_size::<T>(data.len())`).
///
/// `T` need not be `Copy`, so a partially-filled `data` on an early error return would leak (or
/// double-free) the already-written prefix; an internal drop guard drops exactly the completed
/// prefix and nothing else. (Types with `PACK > 1` are `Copy`, so a mid-pack failure that leaves
/// part of a pack written has nothing to drop — see `read_pack`.)
pub(crate) fn read_packed_le_bytes<T: InPlaceCodec>(
    bytes: &[u8],
    data: &mut [MaybeUninit<T>],
) -> Result<(), PrimitiveError> {
    let mut guard = SliceDropGuard::<T>::new(data.as_mut_ptr());
    if T::PACK == 1 {
        for (chunk, slot) in bytes.chunks_exact(T::ENCODED_SIZE).zip(data.iter_mut()) {
            slot.write(T::read_le_bytes(chunk)?);
            guard.inc_len();
        }
    } else {
        let n_full = data.len() / T::PACK;
        let rem = data.len() % T::PACK;
        let (packed, last) = bytes.split_at(n_full * T::PACK_BYTES);
        let mut written = 0usize;
        for chunk in packed.chunks_exact(T::PACK_BYTES) {
            T::read_pack(chunk, &mut data[written..written + T::PACK])?;
            written += T::PACK;
            guard.add_len(T::PACK);
        }
        if rem > 0 {
            if tail_is_unpacked::<T>(rem) {
                // Tail was encoded element-by-element, not padded (see `write_packed_le_bytes`).
                for chunk in last.chunks_exact(T::ENCODED_SIZE) {
                    data[written].write(T::read_le_bytes(chunk)?);
                    written += 1;
                    guard.inc_len();
                }
            } else {
                // The trailing pack was padded with dummy elements up to a full `PACK`-element
                // group on encode (see `write_packed_le_bytes`): decode it whole, then keep only
                // the first `rem` real elements and let the owned `Vec`'s iterator drop the dummy
                // padding.
                let mut padded = Box::<[T]>::new_uninit_slice(T::PACK);
                T::read_pack(last, &mut padded)?;
                // SAFETY: `read_pack` above succeeded, so it initialized every one of the
                // `T::PACK` slots.
                let padded = unsafe { padded.assume_init() };
                for item in Vec::from(padded).into_iter().take(rem) {
                    data[written].write(item);
                    written += 1;
                    guard.inc_len();
                }
            }
        }
    }
    // SAFETY: every slot of `data` was just written above (every `return`/`?` above this point is
    // an error path, so reaching here means the loops ran to completion); forgetting the guard
    // hands the now-fully-initialized `data` back to the caller instead of dropping it as empty.
    std::mem::forget(guard);
    Ok(())
}

/// Drops the initialized prefix of a `*mut MaybeUninit<T>` slice on unwind / early return.
///
/// Used by [`InPlaceCodec`] decoders that fill an uninitialized buffer element by element: if a
/// mid-way decode fails, the already-written prefix must be dropped (and nothing else) to avoid
/// leaks / double-frees for non-`Copy` `T`.
struct SliceDropGuard<T> {
    ptr: *mut MaybeUninit<T>,
    initialized_len: usize,
}

impl<T> SliceDropGuard<T> {
    fn new(ptr: *mut MaybeUninit<T>) -> Self {
        Self {
            ptr,
            initialized_len: 0,
        }
    }

    #[inline(always)]
    #[allow(clippy::arithmetic_side_effects)]
    fn inc_len(&mut self) {
        self.initialized_len += 1;
    }

    #[inline(always)]
    #[allow(clippy::arithmetic_side_effects)]
    fn add_len(&mut self, n: usize) {
        self.initialized_len += n;
    }
}

impl<T> Drop for SliceDropGuard<T> {
    #[inline(always)]
    fn drop(&mut self) {
        unsafe {
            std::ptr::drop_in_place(std::ptr::slice_from_raw_parts_mut(
                self.ptr.cast::<T>(),
                self.initialized_len,
            ));
        }
    }
}