arctic-map 0.1.4

Lock-free adaptive radix tree
Documentation
//! Support for borrowed dynamically sized `&[u8]` keys.

use core::ffi::CStr;
use core::fmt::Debug;
use core::marker::PhantomData;

use ribbit::u13;

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::boxed_slice::BoxedSlice;

/// # Safety
///
/// Implementer must guarantee that `Raw` is unsized
/// and repr(transparent) with `[u8]`.
pub unsafe trait Raw: 'static + AsRef<[u8]> + Debug {
    #[expect(clippy::wrong_self_convention)]
    fn into_boxed(&self) -> Box<Self>;
}
unsafe impl Raw for [u8] {
    #[inline]
    fn into_boxed(&self) -> Box<Self> {
        Box::from(self)
    }
}
unsafe impl Raw for str {
    #[inline]
    fn into_boxed(&self) -> Box<Self> {
        Box::from(self)
    }
}

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

impl<I, R> Slice<I, R>
where
    I: r#unsized::Invariant,
    R: ?Sized + Raw,
{
    /// Construct a slice after validating.
    ///
    /// Returns `Ok` if the input `key` satisfies the invariant.
    pub fn new(key: &R) -> Result<&Self, I::Error> {
        I::validate(key.as_ref())?;
        // Invariants checked above
        Ok(unsafe { Self::new_unchecked(key) })
    }
}

impl<I, R: ?Sized> Slice<I, R> {
    /// # Safety
    ///
    /// Caller must ensure `key` upholds invariants, i.e., `I::validate(key)` would return `Ok`.
    #[inline]
    pub const unsafe fn new_unchecked(key: &R) -> &Self {
        unsafe { core::mem::transmute::<&R, &Self>(key) }
    }

    /// Get a reference to the underlying buffer.
    #[doc(hidden)]
    #[inline]
    pub const fn as_raw(&self) -> &R {
        &self.raw
    }
}

impl<I> Slice<I, str> {
    /// Get a reference to the underlying `str`.
    #[inline]
    pub const fn as_str(&self) -> &str {
        self.as_raw()
    }
}

impl<I> Slice<I, [u8]> {
    /// Get a reference to the underlying `[u8]`.
    #[inline]
    pub const fn as_bytes(&self) -> &[u8] {
        self.as_raw()
    }
}

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

impl<I, R> ToOwned for Slice<I, R>
where
    R: ?Sized + Raw,
{
    type Owned = BoxedSlice<I, R>;
    fn to_owned(&self) -> Self::Owned {
        unsafe { BoxedSlice::new_unchecked(self.as_raw().into_boxed()) }
    }
}

impl<'a> From<&'a CStr> for &'a Slice<Terminated<0>, [u8]> {
    fn from(str: &'a CStr) -> Self {
        // SAFETY: `CStr` is null terminated
        unsafe { Slice::new_unchecked(str.to_bytes_with_nul()) }
    }
}

impl<'a, I, R> crate::Key for &'a Slice<I, R>
where
    I: r#unsized::Invariant,
    R: ?Sized + Raw,
{
    type Borrowed = Slice<I, R>;

    type Insert<'k>
        = &'a Slice<I, R>
    where
        Self: 'k;

    type Read<'k> = Reader<'k, I::Terminate>;
    type Write = Writer<I>;
    type Edge = edge::Slice<I::Terminate>;
    type Len = Byte;

    fn as_insert(&self) -> Self::Insert<'_> {
        self
    }

    fn insert_as_read<'k>(insert: Self::Insert<'k>) -> Self::Read<'k>
    where
        Self: 'k,
    {
        Self::Read::from(insert)
    }

    fn insert_to_key<'k>(insert: Self::Insert<'k>) -> Self
    where
        Self: 'k,
    {
        insert
    }

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

#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
pub struct Reader<'k, T>(pub(crate) r#unsized::boxed_slice::Reader<'k, T>);

impl<'k, I, R> From<&'k Slice<I, R>> for Reader<'k, I::Terminate>
where
    I: r#unsized::Invariant,
    R: ?Sized + Raw,
{
    #[inline]
    fn from(key: &'k Slice<I, R>) -> Self {
        Self(r#unsized::boxed_slice::Reader::from(key))
    }
}

impl<'k, T: Terminate> From<&'k [u8]> for Reader<'k, T> {
    #[inline]
    fn from(prefix: &'k [u8]) -> Self {
        Self(r#unsized::boxed_slice::Reader::from(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())
    }
}

impl<T: Terminate> key::Read for Reader<'_, T> {
    const LEN: Option<Byte> = None;

    type Edge = edge::Slice<T>;
    type Len = Byte;

    fn len(&self) -> Self::Len {
        self.0.len()
    }

    fn get_edge(
        &self,
        len: <ribbit::Packed<Self::Edge> as edge::Meta>::Len,
    ) -> ribbit::Packed<Self::Edge> {
        let min = len.bytes().min(self.0.slice.len());
        edge::Slice::new(&self.0.slice[..min]).with_terminate(T::new(
            self.0.terminate.get() && len.bytes() > self.0.slice.len(),
        ))
    }

    fn get_byte(&self, index: u13) -> Option<u8> {
        self.0.get_byte(index.bytes())
    }

    fn match_prefix(&self, meta: ribbit::Packed<edge::Slice<T>>) -> Self::Len {
        let other = unsafe { meta.as_slice() };

        let index = r#unsized::common_prefix(self.0.slice, other);
        let terminate = self.0.terminate.get()
            && index == self.0.slice.len()
            && index == other.len()
            && meta.terminate().get();

        Byte(index + terminate as usize)
    }

    #[inline]
    fn prefix(self, end: Byte) -> Self {
        Self(self.0.prefix(end))
    }

    #[inline]
    fn suffix(self, start: Byte) -> Self {
        Self(self.0.suffix(start))
    }

    #[inline]
    fn common_prefix(self, other: Self) -> Self {
        Self(self.0.common_prefix(other.0))
    }
}

#[doc(hidden)]
#[derive(Clone, Default, Debug)]
pub struct Writer<I: r#unsized::Invariant> {
    last: ribbit::Packed<edge::Slice<I::Terminate>>,
    len: Byte,
}

impl<I: r#unsized::Invariant> Writer<I> {
    unsafe fn as_slice_unchecked<'a, R: ?Sized>(&self) -> &'a Slice<I, R> {
        let len = self.len.bytes();
        let suffix = unsafe { self.last.as_slice() };
        let raw = I::Terminate::trim(unsafe {
            core::slice::from_raw_parts(suffix.as_ptr().byte_sub(len - suffix.len()), len)
        });
        unsafe { Slice::<I, R>::new_unchecked(core::mem::transmute_copy::<&[u8], &R>(&raw)) }
    }
}

impl<I: r#unsized::Invariant> key::Write<Reader<'_, I::Terminate>> for Writer<I> {
    type Len = Byte;

    fn new(
        prefix: Reader<'_, I::Terminate>,
        key: ribbit::Packed<edge::Slice<I::Terminate>>,
    ) -> (Self, Self::Len) {
        let len = prefix.len() + key.len().into();
        (Writer { last: key, len }, len)
    }

    fn replace(
        &mut self,
        start: Self::Len,
        _: u8,
        edge: ribbit::Packed<edge::Slice<I::Terminate>>,
    ) -> Self::Len {
        validate!(start <= self.len);
        self.len = start + Byte::BYTE + edge.len().into();
        self.last = edge;
        self.len
    }
}