tempest-kv 0.0.2

//! This module contains base types for the key-value storage layer.
//!
//! - [`SeqNum`]: A sequence number new-type, since they are 56 bits, due to bit-packing.
//! - [`KeyKind`]: A one-byte key-kind that identifies different operations in our storage layer.
//! - [`KeyTrailer`]: This type packs the [`SeqNum`] in the higher 7 bytes, with [`KeyKind`] at the
//!   lowest byte. Ordering by it therefore prioritizes the sequence number over the kind.
//! - [`InternalKey`]: This is the key representation that our storage layer uses. It implements
//!   custom ordering through the [`Comparer`] trait, allowing for key suffixes that encode custom
//!   data. Tempest uses the suffix for `HlcTimestamps`, to allow ordering of data across storages.

use std::{cmp, marker::PhantomData};

use bincode::Options as BincodeOptions;
use bytes::Bytes;
use derive_more::Display;
use nonmax::NonMaxU64;
use num_enum::{IntoPrimitive, TryFromPrimitive};
use serde::{Deserialize, Serialize};
use zerocopy::{Immutable, IntoBytes, KnownLayout, LittleEndian, U64};

pub mod comparer;

pub use comparer::*;

use crate::StorageError;

/// The crate wide used [`bincode`] encoding options.
#[doc(hidden)]
pub fn bincode_options() -> impl BincodeOptions {
    bincode::options()
        .with_fixint_encoding() // Important: no variable length ints
        .with_little_endian() // Ensure consistency across platforms
        .allow_trailing_bytes()
}

#[derive(
    derive_more::Debug,
    Display,
    Clone,
    Copy,
    PartialEq,
    Eq,
    PartialOrd,
    Ord,
    Hash,
    Serialize,
    Deserialize,
)]
#[debug("SeqNum({_0:?})")]
pub(crate) struct SeqNum(NonMaxU64);

impl SeqNum {
    /// The zero seqnum represent the absence of any resonable seqnum, used as the lowest bound.
    pub const ZERO: Self = unsafe { Self::new_unchecked(0) };

    /// The sequence number one below the start.
    /// Sometimes exists when searching in a silo, that had no inserts yet.
    pub const MIN: Self = unsafe { Self::new_unchecked(15) };
    /// The first sequence number used for keys.
    /// 1-14 are reserved for later use
    pub const START: Self = unsafe { Self::new_unchecked(16) };
    /// The maximum sequence number available. Sequence numbers may only use 56 bits,
    /// since the [`KeyTrailer`] encodes the [`SeqNum`] in the upper 7 bytes, with the
    /// [`KeyKind`] in the lowest byte, to achieve the best bit packing.
    pub const MAX: Self = unsafe { Self::new_unchecked((1 << 56) - 1) };

    #[cfg(test)]
    pub(crate) const TEST: Self = Self::ZERO;

    pub(crate) const fn new(val: u64) -> Option<Self> {
        if val > Self::MAX.get() {
            return None;
        }
        // SAFETY: Just checked `val` is in valid range
        Some(unsafe { SeqNum::new_unchecked(val) })
    }

    /// Creates a new `SeqNum` without verifying that it is within bounds.
    ///
    /// # Safety
    ///
    /// Caller has to ensure that `val` is at most [`SeqNum::MAX`].
    #[inline]
    pub(crate) const unsafe fn new_unchecked(val: u64) -> Self {
        // SAFETY: User has to ensure that `val <= Self::MAX`
        Self(unsafe { NonMaxU64::new_unchecked(val) })
    }

    // Returns the value as a u64 primitive type.
    #[inline]
    pub(crate) const fn get(&self) -> u64 {
        self.0.get()
    }
}

impl Default for SeqNum {
    fn default() -> Self {
        Self::START
    }
}

// These values are part of the file format and shall never be changed.
#[repr(u8)]
#[derive(
    Debug,
    Display,
    Clone,
    Copy,
    PartialEq,
    Eq,
    PartialOrd,
    Ord,
    Hash,
    IntoPrimitive,
    TryFromPrimitive,
)]
pub enum KeyKind {
    Delete = 0,
    Put = 1,
}

impl KeyKind {
    /// The lowest value (byte-wise).
    pub const MIN: Self = Self::Delete;
    /// The highest value (byte-wise).
    pub const MAX: Self = Self::Put;

    #[cfg(test)]
    pub(crate) const TEST: Self = Self::MIN;
}

#[derive(
    derive_more::Debug,
    Clone,
    Copy,
    PartialEq,
    Eq,
    PartialOrd,
    Ord,
    Hash,
    IntoBytes,
    KnownLayout,
    Immutable,
)]
#[repr(transparent)]
#[debug("KeyTrailer({}: seqnum={:?}, kind={})", _0, self.seqnum(), self.kind())]
pub(crate) struct KeyTrailer(U64<LittleEndian>);

impl KeyTrailer {
    pub(crate) const MAX: Self = Self::new(SeqNum::MAX, KeyKind::MAX);

    #[cfg(test)]
    pub(crate) const TEST: Self = Self::new(SeqNum::TEST, KeyKind::TEST);

    pub(crate) const fn new(seqnum: SeqNum, kind: KeyKind) -> Self {
        Self(U64::new((seqnum.get() << 8) | (kind as u64)))
    }

    pub(crate) const fn seqnum(&self) -> SeqNum {
        // SAFETY: we right shift by 8, so it's less than or equal to SeqNum::MAX
        unsafe { SeqNum::new_unchecked(self.0.get() >> 8) }
    }

    pub(crate) const fn kind(&self) -> KeyKind {
        // SAFETY: we mask by 0xFF to get the key kind bits.
        // The result is always a valid KeyKind because KeyTrailer has no FromBytes impl;
        // all instances go through new() or TryFrom<[u8; 8]>, both of which validate
        // the KeyKind byte.
        unsafe { std::mem::transmute((self.0.get() & 0xFF) as u8) }
    }
}

impl TryFrom<[u8; 8]> for KeyTrailer {
    type Error = StorageError;

    fn try_from(bytes: [u8; 8]) -> Result<Self, Self::Error> {
        let val = u64::from_le_bytes(bytes);
        let kind_byte = (val & 0xFF) as u8;
        let _ = KeyKind::try_from(kind_byte)?;
        Ok(Self(U64::new(val)))
    }
}

#[derive(derive_more::Debug, Clone)]
#[debug("InternalKey(key={}, seqnum={:?}, kind={:?})", C::format(key.as_ref()), trailer.seqnum(), trailer.kind())]
pub struct InternalKey<C: Comparer = DefaultComparer, K: AsRef<[u8]> = Bytes> {
    key: K,
    trailer: KeyTrailer,
    _marker: PhantomData<C>,
}

impl<C: Comparer, K: AsRef<[u8]>> InternalKey<C, K> {
    pub(crate) fn new(key: K, trailer: KeyTrailer) -> Self {
        Self {
            key,
            trailer,
            _marker: PhantomData,
        }
    }

    pub(crate) const fn trailer(&self) -> KeyTrailer {
        self.trailer
    }

    pub const fn key(&self) -> &K {
        &self.key
    }

    pub(crate) fn compare_logical(
        &self,
        other: &InternalKey<C, impl AsRef<[u8]>>,
    ) -> cmp::Ordering {
        C::compare_logical(self.key().as_ref(), other.key().as_ref())
    }

    /// Converts this key to an `InternalKey<C, &[u8]>`, by slicing the key (borrow).
    pub(crate) fn slice_key(&self) -> InternalKey<C, &[u8]> {
        InternalKey::new(self.key.as_ref(), self.trailer)
    }

    /// Converts this key to an `InternalKey<C, Bytes>`, by copying the key (clone).
    pub(crate) fn byte_key(&self) -> InternalKey<C, Bytes> {
        InternalKey::new(Bytes::copy_from_slice(self.key.as_ref()), self.trailer)
    }

    /// Split up this key into `(prefix, suffix)` based on the comparer.
    pub fn split_up(&self) -> (&[u8], &[u8]) {
        let key = self.key.as_ref();
        C::split_up(key)
    }
}

impl<C: Comparer> InternalKey<C, Bytes> {
    /// Generate an internal key for testing other parts of Tempest.
    /// Test keys with a higher ID will be greater than ones with a lower ID.
    /// This is useful to test correctness of internal ordering.
    #[cfg(test)]
    pub(crate) fn test(id: u64) -> Self {
        let test_key = Bytes::from(id.to_be_bytes().to_vec());
        Self::new(test_key, KeyTrailer::TEST)
    }

    #[cfg(test)]
    pub(crate) fn test_with_seqnum(id: u64, seqnum: u64) -> Self {
        let test_key = Bytes::from(id.to_be_bytes().to_vec());
        Self::new(
            test_key,
            KeyTrailer::new(SeqNum::new(seqnum).unwrap(), KeyKind::TEST),
        )
    }

    // Test helper that gets the key as an easily comparable `u64` value.
    #[cfg(test)]
    pub(crate) fn test_key_as_u64(&self) -> u64 {
        let mut buf = [0u8; 8];
        buf.copy_from_slice(self.key.as_ref());
        u64::from_be_bytes(buf)
    }
}

impl<C: Comparer, K: AsRef<[u8]>> cmp::PartialEq for InternalKey<C, K> {
    fn eq(&self, other: &Self) -> bool {
        self.cmp(other).is_eq()
    }
}

impl<C: Comparer, K: AsRef<[u8]>> cmp::Eq for InternalKey<C, K> {}

impl<C: Comparer, K: AsRef<[u8]>> cmp::PartialOrd for InternalKey<C, K> {
    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl<C: Comparer, K: AsRef<[u8]>> cmp::Ord for InternalKey<C, K> {
    fn cmp(&self, other: &Self) -> cmp::Ordering {
        match C::compare_physical(self.key.as_ref(), other.key.as_ref()) {
            cmp::Ordering::Equal => other.trailer.cmp(&self.trailer),
            ord => ord,
        }
    }
}