armdb 0.2.0

sharded bitcask key-value storage optimized for NVMe
Documentation
use crate::error::DbResult;

/// Codec for encoding/decoding typed values to/from bytes.
///
/// Implemented as zero-sized marker structs for zero runtime overhead.
/// Used by [`TypedTree`](crate::TypedTree) for disk serialization.
pub trait Codec<T>: Send + Sync {
    /// Encode `value` into `buf`. Implementations must clear `buf` before writing —
    /// the caller may pass a reusable buffer.
    fn encode_to(&self, value: &T, buf: &mut Vec<u8>) -> DbResult<()>;
    fn decode_from(&self, bytes: &[u8]) -> DbResult<T>;
    /// Encoded byte size of `value`. Default: encode into a throwaway buffer
    /// and return the resulting length. Specialize for cheap-size codecs.
    fn size(&self, value: &T) -> usize {
        let mut buf = Vec::new();
        let _ = self.encode_to(value, &mut buf);
        buf.len()
    }
}

#[cfg(feature = "rapira-codec")]
#[derive(Default, Clone, Debug)]
pub struct RapiraCodec;

#[cfg(feature = "rapira-codec")]
impl<T: rapira::Rapira> Codec<T> for RapiraCodec {
    fn encode_to(&self, value: &T, buf: &mut Vec<u8>) -> DbResult<()> {
        buf.clear();
        rapira::extend_vec(value, buf);
        Ok(())
    }

    fn decode_from(&self, bytes: &[u8]) -> DbResult<T> {
        unsafe {
            rapira::deser_unchecked(bytes)
                .map_err(|_| crate::error::DbError::CorruptedEntry { offset: 0 })
        }
    }

    fn size(&self, value: &T) -> usize {
        rapira::Rapira::size(value)
    }
}

#[cfg(feature = "bitcode-codec")]
#[derive(Default, Clone, Debug)]
pub struct BitcodeCodec;

#[cfg(feature = "bitcode-codec")]
impl<T: bitcode::Encode + for<'a> bitcode::Decode<'a>> Codec<T> for BitcodeCodec {
    fn encode_to(&self, value: &T, buf: &mut Vec<u8>) -> DbResult<()> {
        buf.clear();
        buf.extend_from_slice(&bitcode::encode(value));
        Ok(())
    }

    fn decode_from(&self, bytes: &[u8]) -> DbResult<T> {
        bitcode::decode(bytes).map_err(|_| crate::error::DbError::CorruptedEntry { offset: 0 })
    }
}

/// Codec using [postcard](https://docs.rs/postcard) — compact binary serde format
/// with varint encoding. Good JS/TS interop via `postcard-js` npm package.
#[cfg(feature = "postcard-codec")]
#[derive(Default, Clone, Debug)]
pub struct PostcardCodec;

#[cfg(feature = "postcard-codec")]
impl<T: serde::Serialize + for<'a> serde::Deserialize<'a>> Codec<T> for PostcardCodec {
    fn encode_to(&self, value: &T, buf: &mut Vec<u8>) -> DbResult<()> {
        buf.clear();
        let encoded = postcard::to_allocvec(value)
            .map_err(|_| crate::error::DbError::Client("postcard serialization failed"))?;
        buf.extend_from_slice(&encoded);
        Ok(())
    }

    fn decode_from(&self, bytes: &[u8]) -> DbResult<T> {
        postcard::from_bytes(bytes).map_err(|_| crate::error::DbError::CorruptedEntry { offset: 0 })
    }
}

/// Codec for `T: IntoBytes + FromBytes` — fixed-size values via zerocopy.
///
/// Zero-cost conversion via `zerocopy::IntoBytes::as_bytes` / `FromBytes::read_from_bytes`.
/// Unlike [`BytemuckCodec`], does not require `bytemuck` — uses the always-available
/// `zerocopy` crate.
#[derive(Default, Clone, Debug)]
pub struct ZerocopyCodec;

impl<T> Codec<T> for ZerocopyCodec
where
    T: zerocopy::IntoBytes + zerocopy::FromBytes + zerocopy::Immutable,
{
    fn encode_to(&self, value: &T, buf: &mut Vec<u8>) -> DbResult<()> {
        buf.clear();
        buf.extend_from_slice(zerocopy::IntoBytes::as_bytes(value));
        Ok(())
    }

    fn decode_from(&self, bytes: &[u8]) -> DbResult<T> {
        zerocopy::FromBytes::read_from_bytes(bytes)
            .map_err(|_| crate::error::DbError::CorruptedEntry { offset: 0 })
    }

    fn size(&self, _value: &T) -> usize {
        size_of::<T>()
    }
}

/// Codec for `T: bytemuck::Pod` — fixed-size values cast as raw bytes.
///
/// Zero-cost conversion via `bytemuck::bytes_of` / `try_from_bytes`.
/// For variable-length vectors use [`BytemuckSliceCodec`] instead.
#[cfg(feature = "bytemuck-codec")]
#[derive(Default, Clone, Debug)]
pub struct BytemuckCodec;

#[cfg(feature = "bytemuck-codec")]
impl<T: bytemuck::Pod> Codec<T> for BytemuckCodec {
    fn encode_to(&self, value: &T, buf: &mut Vec<u8>) -> DbResult<()> {
        buf.clear();
        buf.extend_from_slice(bytemuck::bytes_of(value));
        Ok(())
    }

    fn decode_from(&self, bytes: &[u8]) -> DbResult<T> {
        if bytes.len() != size_of::<T>() {
            return Err(crate::error::DbError::CorruptedEntry { offset: 0 });
        }
        Ok(bytemuck::pod_read_unaligned(bytes))
    }

    fn size(&self, _value: &T) -> usize {
        size_of::<T>()
    }
}

/// Codec for `Vec<T>` where `T: bytemuck::Pod` — variable-length vector
/// serialized as contiguous raw bytes via `bytemuck::cast_slice`.
///
/// Matches the pattern from `bytemuck_slice_record!` in the `armour` crate.
///
/// ```ignore
/// let tree = TypedTree::<16, Vec<u64>, BytemuckSliceCodec>::open(config, BytemuckSliceCodec)?;
/// tree.put(&key, vec![1u64, 2, 3])?;
/// ```
#[cfg(feature = "bytemuck-codec")]
#[derive(Default, Clone, Debug)]
pub struct BytemuckSliceCodec;

#[cfg(feature = "bytemuck-codec")]
impl<W: BytemuckVec> Codec<W> for BytemuckSliceCodec {
    fn encode_to(&self, value: &W, buf: &mut Vec<u8>) -> DbResult<()> {
        buf.clear();
        buf.extend_from_slice(bytemuck::cast_slice(value.as_slice()));
        Ok(())
    }

    fn decode_from(&self, bytes: &[u8]) -> DbResult<W> {
        let item_size = size_of::<W::Item>();
        if item_size == 0 {
            return Ok(W::from_vec(Vec::new()));
        }
        if !bytes.is_empty() && !bytes.len().is_multiple_of(item_size) {
            return Err(crate::error::DbError::CorruptedEntry { offset: 0 });
        }
        let vec: Vec<W::Item> = bytes
            .chunks_exact(item_size)
            .map(bytemuck::pod_read_unaligned)
            .collect();
        Ok(W::from_vec(vec))
    }

    fn size(&self, value: &W) -> usize {
        std::mem::size_of_val(value.as_slice())
    }
}

/// Trait for types that can be encoded/decoded as a contiguous slice of `Pod` elements.
///
/// Implemented automatically for `Vec<T: Pod>`. For newtype wrappers over `Vec<T>`,
/// implement this trait to use [`BytemuckSliceCodec`] directly:
///
/// ```ignore
/// #[derive(From, AsRef)]
/// #[repr(transparent)]
/// pub struct PostLikes(pub Vec<UserID>);
///
/// impl BytemuckVec for PostLikes {
///     type Item = UserID;
///     fn from_vec(v: Vec<UserID>) -> Self { Self(v) }
///     fn as_slice(&self) -> &[UserID] { &self.0 }
/// }
///
/// let map = db.open_typed_map::<PostLikes, BytemuckSliceCodec>(...)?;
/// ```
#[cfg(feature = "bytemuck-codec")]
pub trait BytemuckVec: Send + Sync + Sized {
    type Item: bytemuck::Pod;
    fn from_vec(v: Vec<Self::Item>) -> Self;
    fn as_slice(&self) -> &[Self::Item];
}

#[cfg(feature = "bytemuck-codec")]
impl<T: bytemuck::Pod + Send + Sync> BytemuckVec for Vec<T> {
    type Item = T;
    fn from_vec(v: Vec<T>) -> Self {
        v
    }
    fn as_slice(&self) -> &[T] {
        self
    }
}

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

    #[test]
    fn zerocopy_size_is_constant() {
        let c = ZerocopyCodec;
        let v: u64 = 42;
        assert_eq!(<ZerocopyCodec as Codec<u64>>::size(&c, &v), 8);
    }

    #[cfg(feature = "rapira-codec")]
    #[test]
    fn rapira_size_matches_encode_len() {
        let c = RapiraCodec;
        let v: Vec<u8> = vec![1, 2, 3, 4, 5];
        let mut buf = Vec::new();
        c.encode_to(&v, &mut buf).unwrap();
        assert_eq!(<RapiraCodec as Codec<Vec<u8>>>::size(&c, &v), buf.len());
    }

    #[cfg(feature = "bytemuck-codec")]
    #[test]
    fn bytemuck_size_is_constant() {
        let c = BytemuckCodec;
        let v: [u8; 16] = [0; 16];
        assert_eq!(<BytemuckCodec as Codec<[u8; 16]>>::size(&c, &v), 16);
    }

    #[cfg(feature = "bytemuck-codec")]
    #[test]
    fn bytemuck_slice_size_matches() {
        let c = BytemuckSliceCodec;
        let v: Vec<u32> = vec![1, 2, 3, 4];
        let mut buf = Vec::new();
        c.encode_to(&v, &mut buf).unwrap();
        assert_eq!(
            <BytemuckSliceCodec as Codec<Vec<u32>>>::size(&c, &v),
            buf.len()
        );
    }
}