arctic-wt 0.1.5

//! Support for owned dynamically sized keys ([`Vec<u8>`], [`Box<[u8]>`][Box]).

use core::borrow::Borrow;
use core::fmt::Debug;
use core::marker::PhantomData;
use core::ops::Deref;
use core::ptr::NonNull;
use std::ffi::CString;

#[cfg(feature = "proptest")]
use proptest::prelude::Strategy;
use ribbit::u6;

use crate::Key;
#[cfg(feature = "proptest")]
use crate::key::Invariant;
use crate::key::Terminated;
use crate::raw::edge;
use crate::raw::edge::Len as _;
use crate::raw::edge::Meta as _;
use crate::raw::key;
use crate::raw::key::Byte;
use crate::raw::key::Len as _;
use crate::raw::key::Read as _;
use crate::raw::key::r#unsized;
use crate::raw::key::r#unsized::Terminate;
use crate::raw::key::r#unsized::slice::Slice;

/// An owned, dynamically sized key that satisfies an [`Invariant`][crate::key::unsized::Invariant].
#[repr(transparent)]
#[derive(Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
pub struct BoxedSlice<I, R: ?Sized = [u8]> {
    invariant: PhantomData<I>,
    raw: Box<R>,
}

impl<I, R: ?Sized> Clone for BoxedSlice<I, R>
where
    Box<R>: Clone,
{
    #[inline]
    fn clone(&self) -> Self {
        Self {
            invariant: PhantomData,
            raw: self.raw.clone(),
        }
    }
}

impl<I, R: ?Sized> Default for BoxedSlice<I, R>
where
    Box<R>: Default,
{
    fn default() -> Self {
        Self {
            invariant: PhantomData,
            raw: Default::default(),
        }
    }
}

impl<I, R> BoxedSlice<I, R>
where
    I: r#unsized::Invariant,
    R: ?Sized + r#unsized::slice::Raw,
{
    /// Construct a boxed slice after validating.
    ///
    /// Returns `Ok` if the input boxed slice satisfies the invariant.
    #[inline]
    pub fn new(key: impl Into<Box<R>>) -> Result<Self, (Box<R>, I::Error)> {
        let key = key.into();
        match Slice::<I, R>::new(&key) {
            Ok(_) => Ok(unsafe { Self::new_unchecked(key) }),
            Err(error) => Err((key, error)),
        }
    }
}

impl<I, R: ?Sized> BoxedSlice<I, R> {
    /// Construct a boxed slice without validating.
    ///
    /// # SAFETY
    ///
    /// Caller must guarantee that `raw` satisfies the invariant, i.e.,
    /// `I::validate(key)` would return `Ok(())`.
    #[inline]
    pub const unsafe fn new_unchecked(key: Box<R>) -> Self {
        Self {
            invariant: PhantomData,
            raw: key,
        }
    }

    /// Get a borrowed [`Slice`] that preserves the invariant.
    #[inline]
    pub const fn as_slice(&self) -> &Slice<I, R> {
        unsafe { Slice::new_unchecked(&self.raw) }
    }

    /// Get an owned boxed slice.
    #[inline]
    pub fn into_boxed_slice(self) -> Box<R> {
        self.raw
    }
}

impl<I, R: ?Sized> Deref for BoxedSlice<I, R> {
    type Target = Slice<I, R>;
    #[inline]
    fn deref(&self) -> &Self::Target {
        self.as_slice()
    }
}

impl<I, R: ?Sized> Borrow<Slice<I, R>> for BoxedSlice<I, R> {
    #[inline]
    fn borrow(&self) -> &Slice<I, R> {
        self.as_slice()
    }
}

impl<I, R: ?Sized> AsRef<Slice<I, R>> for BoxedSlice<I, R> {
    #[inline]
    fn as_ref(&self) -> &Slice<I, R> {
        self.as_slice()
    }
}

impl From<CString> for BoxedSlice<Terminated<0>, [u8]> {
    fn from(string: CString) -> Self {
        // SAFETY: `CString` is null terminated
        unsafe { Self::new_unchecked(string.into_bytes_with_nul().into_boxed_slice()) }
    }
}

#[cfg(feature = "proptest")]
impl<I, R> proptest::arbitrary::Arbitrary for BoxedSlice<I, R>
where
    I: Invariant,
    R: ?Sized + r#unsized::slice::Raw + core::fmt::Debug,
    Box<R>: proptest::arbitrary::Arbitrary,
{
    type Parameters = <Box<R> as proptest::arbitrary::Arbitrary>::Parameters;
    type Strategy = proptest::strategy::BoxedStrategy<Self>;

    fn arbitrary_with(args: Self::Parameters) -> Self::Strategy {
        <Box<R>>::arbitrary_with(args)
            .prop_filter_map("Invariant violated", |boxed_slice| {
                Self::new(boxed_slice).ok()
            })
            .boxed()
    }
}

#[cfg(feature = "rand")]
impl rand::distr::Distribution<BoxedSlice<r#unsized::NonNull, str>>
    for rand::distr::StandardUniform
{
    fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> BoxedSlice<r#unsized::NonNull, str> {
        let uniform = rand::distr::Uniform::new_inclusive(1 as char, char::MAX).unwrap();
        let string = rand::distr::SampleString::sample_string(&uniform, rng, 32);
        unsafe { BoxedSlice::new_unchecked(string.into_boxed_str()) }
    }
}

impl<I, R> Key for BoxedSlice<I, R>
where
    I: r#unsized::Invariant,
    R: ?Sized + r#unsized::slice::Raw,
{
    type Read<'k> = Reader<'k, I::Terminate>;
    type Write = Writer;
    type Borrowed = Slice<I, R>;
    type Insert<'k> = &'k Slice<I, R>;
    type Edge = edge::Le;
    type Len = Byte;

    #[inline]
    fn as_insert(&self) -> Self::Insert<'_> {
        self.as_slice()
    }

    #[inline]
    fn insert_as_read<'k>(insert: Self::Insert<'k>) -> Self::Read<'k>
    where
        Self: 'k,
    {
        Reader::from(insert)
    }

    #[inline]
    fn insert_to_key<'k>(insert: Self::Insert<'k>) -> Self
    where
        Self: 'k,
    {
        insert.to_owned()
    }

    #[inline]
    unsafe fn write_as_insert<'k>(writer: &'k Self::Write) -> Self::Insert<'k>
    where
        Self: 'k,
    {
        unsafe { writer.as_slice_unchecked() }
    }
}

impl<'k, I, R> From<&'k Slice<I, R>> for Reader<'k, I::Terminate>
where
    I: r#unsized::Invariant,
    R: ?Sized + r#unsized::slice::Raw,
{
    #[inline]
    fn from(slice: &'k Slice<I, R>) -> Self {
        let slice = slice.as_raw().as_ref();
        Self {
            terminate: I::Terminate::TRUE,
            ..Self::new_prefix(slice)
        }
    }
}

impl<'k, T: Terminate> From<&'k [u8]> for Reader<'k, T> {
    #[inline]
    fn from(prefix: &'k [u8]) -> Self {
        Reader::new_prefix(prefix)
    }
}

impl<'k, T: Terminate> From<&'k str> for Reader<'k, T> {
    #[inline]
    fn from(prefix: &'k str) -> Self {
        Self::from(prefix.as_bytes())
    }
}

impl<'k, const N: usize, T: Terminate> From<&'k [u8; N]> for Reader<'k, T> {
    #[inline]
    fn from(prefix: &'k [u8; N]) -> Self {
        Self::from(prefix.as_slice())
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct Reader<'k, T> {
    // Not using slices in order to preserve the
    // pointer provenance of the original slice for
    // `crate::raw::key::slice::Writer` and `crate::raw::edge::Slice`.
    ptr: NonNull<u8>,
    pub(crate) len: usize,
    pub(super) terminate: T,
    slice: PhantomData<&'k [u8]>,
}

impl<'k, T: Default> Reader<'k, T> {
    #[inline]
    pub(crate) fn new_prefix(prefix: &'k [u8]) -> Self {
        Self {
            ptr: NonNull::from(prefix).cast::<u8>(),
            len: prefix.len(),
            terminate: T::default(),
            slice: PhantomData,
        }
    }
}

#[expect(private_bounds)]
impl<'k, T: Terminate> Reader<'k, T> {
    #[inline]
    pub(crate) fn get_byte(&self, index: usize) -> Option<u8> {
        if let Some(byte) = self.as_slice().get(index) {
            return Some(*byte);
        }

        (self.terminate.get() && index == self.len).then_some(0)
    }

    #[inline]
    pub(crate) fn as_slice(&self) -> &'k [u8] {
        // SAFETY: `self.ptr` and `self.len` form a valid slice for lifetime 'k
        unsafe { &*core::ptr::slice_from_raw_parts(self.ptr.as_ptr().cast_const(), self.len) }
    }

    #[inline]
    pub(super) fn as_non_null(&self) -> NonNull<u8> {
        self.ptr
    }
}

impl<T: Default> Default for Reader<'_, T> {
    #[inline]
    fn default() -> Self {
        Self::new_prefix(&[])
    }
}

impl<T: Terminate> key::Read for Reader<'_, T> {
    const LEN: Option<Self::Len> = None;
    type Edge = edge::Le;
    type Len = Byte;

    #[inline]
    fn len(&self) -> Self::Len {
        Byte(self.len + self.terminate.get() as usize)
    }

    #[inline]
    fn get_edge(
        &self,
        len: <ribbit::Packed<Self::Edge> as edge::Meta>::Len,
    ) -> ribbit::Packed<Self::Edge> {
        let len = u6::new((self.len().bits()).min(len.bits()) as u8);
        edge::Le::new(r#unsized::read_u64(self.as_slice()), len)
    }

    #[inline]
    fn get_byte(&self, index: u6) -> Option<u8> {
        self.get_byte(index.bytes())
    }

    #[inline]
    fn match_exact(
        &self,
        edge: <Self::Edge as ribbit::Pack>::Packed,
    ) -> Option<<ribbit::Packed<Self::Edge> as edge::Meta>::Len> {
        // Avoid bit <-> byte conversion
        let len_edge = edge.len();
        let len_match = (edge.raw() ^ r#unsized::read_u64(self.as_slice())).trailing_zeros() as u8;
        (len_match >= len_edge.value()).then_some(len_edge)
    }

    #[inline]
    fn match_prefix(&self, edge: <Self::Edge as ribbit::Pack>::Packed) -> Self::Len {
        Byte(((edge.raw() ^ r#unsized::read_u64(self.as_slice())).trailing_zeros() as usize) >> 3)
    }

    #[inline]
    fn prefix(self, end: Self::Len) -> Self {
        validate!(end <= self.len());
        let end = end.bytes();

        Self {
            ptr: self.ptr,
            len: self.len.min(end),
            terminate: T::new(self.terminate.get() && (end > self.len)),
            slice: PhantomData,
        }
    }

    #[inline]
    fn suffix(self, start: Self::Len) -> Self {
        validate!(start <= self.len());
        let start = start.bytes();
        let offset = self.len.min(start);

        Self {
            len: self.len - offset,
            // NOTE: slice key implementation requires us to preserve the
            // `self.slice` pointer, even if the slice is empty.
            ptr: unsafe { self.ptr.byte_add(offset) },
            terminate: T::new(self.terminate.get() && (start <= self.len)),
            slice: PhantomData,
        }
    }

    #[inline]
    fn common_prefix(self, other: Self) -> Self {
        let index = r#unsized::common_prefix(self.as_slice(), other.as_slice());

        Self {
            ptr: self.ptr,
            len: index,
            terminate: T::new(
                self.terminate.get()
                    && other.terminate.get()
                    && index == self.len
                    && index == other.len,
            ),
            slice: PhantomData,
        }
    }
}

#[doc(hidden)]
#[repr(transparent)]
#[derive(Debug, Default)]
pub struct Writer(Vec<u8>);

impl Writer {
    unsafe fn as_slice_unchecked<I: r#unsized::Invariant, R: ?Sized>(&self) -> &Slice<I, R> {
        let raw = I::Terminate::trim(self.0.as_slice());
        unsafe { Slice::<I, R>::new_unchecked(core::mem::transmute_copy::<&[u8], &R>(&raw)) }
    }
}

impl<'k, T: Terminate> key::Write<Reader<'k, T>> for Writer {
    type Len = Byte;

    #[inline]
    fn new(prefix: Reader<'k, T>, key: ribbit::Packed<edge::Le>) -> (Self, Self::Len) {
        let len = prefix.len() + key.len().into();
        let mut buffer = Vec::new();
        buffer.extend_from_slice(prefix.as_slice());
        if prefix.terminate.get() {
            buffer.push(u8::MIN);
            validate_eq!(key.len().bits(), 0);
        } else {
            buffer.extend(key);
        }
        (Writer(buffer), len)
    }

    #[inline]
    fn replace(&mut self, start: Self::Len, node: u8, edge: ribbit::Packed<edge::Le>) -> Self::Len {
        validate!(start.0 <= self.0.len());
        self.0.truncate(start.0);
        self.0.push(node);
        self.0.extend(edge);
        Byte(self.0.len())
    }
}