Skip to main content

arctic/raw/key/unsized/
boxed_slice.rs

1//! Support for owned dynamically sized keys ([`Vec<u8>`], [`Box<[u8]>`][Box]).
2
3use core::borrow::Borrow;
4use core::fmt::Debug;
5use core::marker::PhantomData;
6use core::ops::Deref;
7use core::ptr::NonNull;
8use std::ffi::CString;
9
10#[cfg(feature = "proptest")]
11use proptest::prelude::Strategy;
12use ribbit::u6;
13
14use crate::Key;
15#[cfg(feature = "proptest")]
16use crate::key::Invariant;
17use crate::key::Terminated;
18use crate::raw::edge;
19use crate::raw::edge::Len as _;
20use crate::raw::edge::Meta as _;
21use crate::raw::key;
22use crate::raw::key::Byte;
23use crate::raw::key::Len as _;
24use crate::raw::key::Read as _;
25use crate::raw::key::r#unsized;
26use crate::raw::key::r#unsized::Terminate;
27use crate::raw::key::r#unsized::slice::Slice;
28
29/// An owned, dynamically sized key that satisfies an [`Invariant`][crate::key::unsized::Invariant].
30#[repr(transparent)]
31#[derive(Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
32pub struct BoxedSlice<I, R: ?Sized = [u8]> {
33    invariant: PhantomData<I>,
34    raw: Box<R>,
35}
36
37impl<I, R: ?Sized> Clone for BoxedSlice<I, R>
38where
39    Box<R>: Clone,
40{
41    #[inline]
42    fn clone(&self) -> Self {
43        Self {
44            invariant: PhantomData,
45            raw: self.raw.clone(),
46        }
47    }
48}
49
50impl<I, R: ?Sized> Default for BoxedSlice<I, R>
51where
52    Box<R>: Default,
53{
54    fn default() -> Self {
55        Self {
56            invariant: PhantomData,
57            raw: Default::default(),
58        }
59    }
60}
61
62impl<I, R> BoxedSlice<I, R>
63where
64    I: r#unsized::Invariant,
65    R: ?Sized + r#unsized::slice::Raw,
66{
67    /// Construct a boxed slice after validating.
68    ///
69    /// Returns `Ok` if the input boxed slice satisfies the invariant.
70    #[inline]
71    pub fn new(key: impl Into<Box<R>>) -> Result<Self, (Box<R>, I::Error)> {
72        let key = key.into();
73        match Slice::<I, R>::new(&key) {
74            Ok(_) => Ok(unsafe { Self::new_unchecked(key) }),
75            Err(error) => Err((key, error)),
76        }
77    }
78}
79
80impl<I, R: ?Sized> BoxedSlice<I, R> {
81    /// Construct a boxed slice without validating.
82    ///
83    /// # SAFETY
84    ///
85    /// Caller must guarantee that `raw` satisfies the invariant, i.e.,
86    /// `I::validate(key)` would return `Ok(())`.
87    #[inline]
88    pub const unsafe fn new_unchecked(key: Box<R>) -> Self {
89        Self {
90            invariant: PhantomData,
91            raw: key,
92        }
93    }
94
95    /// Get a borrowed [`Slice`] that preserves the invariant.
96    #[inline]
97    pub const fn as_slice(&self) -> &Slice<I, R> {
98        unsafe { Slice::new_unchecked(&self.raw) }
99    }
100
101    /// Get an owned boxed slice.
102    #[inline]
103    pub fn into_boxed_slice(self) -> Box<R> {
104        self.raw
105    }
106}
107
108impl<I, R: ?Sized> Deref for BoxedSlice<I, R> {
109    type Target = Slice<I, R>;
110    #[inline]
111    fn deref(&self) -> &Self::Target {
112        self.as_slice()
113    }
114}
115
116impl<I, R: ?Sized> Borrow<Slice<I, R>> for BoxedSlice<I, R> {
117    #[inline]
118    fn borrow(&self) -> &Slice<I, R> {
119        self.as_slice()
120    }
121}
122
123impl<I, R: ?Sized> AsRef<Slice<I, R>> for BoxedSlice<I, R> {
124    #[inline]
125    fn as_ref(&self) -> &Slice<I, R> {
126        self.as_slice()
127    }
128}
129
130impl From<CString> for BoxedSlice<Terminated<0>, [u8]> {
131    fn from(string: CString) -> Self {
132        // SAFETY: `CString` is null terminated
133        unsafe { Self::new_unchecked(string.into_bytes_with_nul().into_boxed_slice()) }
134    }
135}
136
137#[cfg(feature = "proptest")]
138impl<I, R> proptest::arbitrary::Arbitrary for BoxedSlice<I, R>
139where
140    I: Invariant,
141    R: ?Sized + r#unsized::slice::Raw + core::fmt::Debug,
142    Box<R>: proptest::arbitrary::Arbitrary,
143{
144    type Parameters = <Box<R> as proptest::arbitrary::Arbitrary>::Parameters;
145    type Strategy = proptest::strategy::BoxedStrategy<Self>;
146
147    fn arbitrary_with(args: Self::Parameters) -> Self::Strategy {
148        <Box<R>>::arbitrary_with(args)
149            .prop_filter_map("Invariant violated", |boxed_slice| {
150                Self::new(boxed_slice).ok()
151            })
152            .boxed()
153    }
154}
155
156#[cfg(feature = "rand")]
157impl rand::distr::Distribution<BoxedSlice<r#unsized::NonNull, str>>
158    for rand::distr::StandardUniform
159{
160    fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> BoxedSlice<r#unsized::NonNull, str> {
161        let uniform = rand::distr::Uniform::new_inclusive(1 as char, char::MAX).unwrap();
162        let string = rand::distr::SampleString::sample_string(&uniform, rng, 32);
163        unsafe { BoxedSlice::new_unchecked(string.into_boxed_str()) }
164    }
165}
166
167impl<I, R> Key for BoxedSlice<I, R>
168where
169    I: r#unsized::Invariant,
170    R: ?Sized + r#unsized::slice::Raw,
171{
172    type Read<'k> = Reader<'k, I::Terminate>;
173    type Write = Writer;
174    type Borrowed = Slice<I, R>;
175    type Insert<'k> = &'k Slice<I, R>;
176    type Edge = edge::Le;
177    type Len = Byte;
178
179    #[inline]
180    fn as_insert(&self) -> Self::Insert<'_> {
181        self.as_slice()
182    }
183
184    #[inline]
185    fn insert_as_read<'k>(insert: Self::Insert<'k>) -> Self::Read<'k>
186    where
187        Self: 'k,
188    {
189        Reader::from(insert)
190    }
191
192    #[inline]
193    fn insert_to_key<'k>(insert: Self::Insert<'k>) -> Self
194    where
195        Self: 'k,
196    {
197        insert.to_owned()
198    }
199
200    #[inline]
201    unsafe fn write_as_insert<'k>(writer: &'k Self::Write) -> Self::Insert<'k>
202    where
203        Self: 'k,
204    {
205        unsafe { writer.as_slice_unchecked() }
206    }
207}
208
209impl<'k, I, R> From<&'k Slice<I, R>> for Reader<'k, I::Terminate>
210where
211    I: r#unsized::Invariant,
212    R: ?Sized + r#unsized::slice::Raw,
213{
214    #[inline]
215    fn from(slice: &'k Slice<I, R>) -> Self {
216        let slice = slice.as_raw().as_ref();
217        Self {
218            terminate: I::Terminate::TRUE,
219            ..Self::new_prefix(slice)
220        }
221    }
222}
223
224impl<'k, T: Terminate> From<&'k [u8]> for Reader<'k, T> {
225    #[inline]
226    fn from(prefix: &'k [u8]) -> Self {
227        Reader::new_prefix(prefix)
228    }
229}
230
231impl<'k, T: Terminate> From<&'k str> for Reader<'k, T> {
232    #[inline]
233    fn from(prefix: &'k str) -> Self {
234        Self::from(prefix.as_bytes())
235    }
236}
237
238impl<'k, const N: usize, T: Terminate> From<&'k [u8; N]> for Reader<'k, T> {
239    #[inline]
240    fn from(prefix: &'k [u8; N]) -> Self {
241        Self::from(prefix.as_slice())
242    }
243}
244
245#[derive(Copy, Clone, Debug, PartialEq, Eq)]
246pub struct Reader<'k, T> {
247    // Not using slices in order to preserve the
248    // pointer provenance of the original slice for
249    // `crate::raw::key::slice::Writer` and `crate::raw::edge::Slice`.
250    ptr: NonNull<u8>,
251    pub(crate) len: usize,
252    pub(super) terminate: T,
253    slice: PhantomData<&'k [u8]>,
254}
255
256impl<'k, T: Default> Reader<'k, T> {
257    #[inline]
258    pub(crate) fn new_prefix(prefix: &'k [u8]) -> Self {
259        Self {
260            ptr: NonNull::from(prefix).cast::<u8>(),
261            len: prefix.len(),
262            terminate: T::default(),
263            slice: PhantomData,
264        }
265    }
266}
267
268#[expect(private_bounds)]
269impl<'k, T: Terminate> Reader<'k, T> {
270    #[inline]
271    pub(crate) fn get_byte(&self, index: usize) -> Option<u8> {
272        if let Some(byte) = self.as_slice().get(index) {
273            return Some(*byte);
274        }
275
276        (self.terminate.get() && index == self.len).then_some(0)
277    }
278
279    #[inline]
280    pub(crate) fn as_slice(&self) -> &'k [u8] {
281        // SAFETY: `self.ptr` and `self.len` form a valid slice for lifetime 'k
282        unsafe { &*core::ptr::slice_from_raw_parts(self.ptr.as_ptr().cast_const(), self.len) }
283    }
284
285    #[inline]
286    pub(super) fn as_non_null(&self) -> NonNull<u8> {
287        self.ptr
288    }
289}
290
291impl<T: Default> Default for Reader<'_, T> {
292    #[inline]
293    fn default() -> Self {
294        Self::new_prefix(&[])
295    }
296}
297
298impl<T: Terminate> key::Read for Reader<'_, T> {
299    const LEN: Option<Self::Len> = None;
300    type Edge = edge::Le;
301    type Len = Byte;
302
303    #[inline]
304    fn len(&self) -> Self::Len {
305        Byte(self.len + self.terminate.get() as usize)
306    }
307
308    #[inline]
309    fn get_edge(
310        &self,
311        len: <ribbit::Packed<Self::Edge> as edge::Meta>::Len,
312    ) -> ribbit::Packed<Self::Edge> {
313        let len = u6::new((self.len().bits()).min(len.bits()) as u8);
314        edge::Le::new(r#unsized::read_u64(self.as_slice()), len)
315    }
316
317    #[inline]
318    fn get_byte(&self, index: u6) -> Option<u8> {
319        self.get_byte(index.bytes())
320    }
321
322    #[inline]
323    fn match_exact(
324        &self,
325        edge: <Self::Edge as ribbit::Pack>::Packed,
326    ) -> Option<<ribbit::Packed<Self::Edge> as edge::Meta>::Len> {
327        // Avoid bit <-> byte conversion
328        let len_edge = edge.len();
329        let len_match = (edge.raw() ^ r#unsized::read_u64(self.as_slice())).trailing_zeros() as u8;
330        (len_match >= len_edge.value()).then_some(len_edge)
331    }
332
333    #[inline]
334    fn match_prefix(&self, edge: <Self::Edge as ribbit::Pack>::Packed) -> Self::Len {
335        Byte(((edge.raw() ^ r#unsized::read_u64(self.as_slice())).trailing_zeros() as usize) >> 3)
336    }
337
338    #[inline]
339    fn prefix(self, end: Self::Len) -> Self {
340        validate!(end <= self.len());
341        let end = end.bytes();
342
343        Self {
344            ptr: self.ptr,
345            len: self.len.min(end),
346            terminate: T::new(self.terminate.get() && (end > self.len)),
347            slice: PhantomData,
348        }
349    }
350
351    #[inline]
352    fn suffix(self, start: Self::Len) -> Self {
353        validate!(start <= self.len());
354        let start = start.bytes();
355        let offset = self.len.min(start);
356
357        Self {
358            len: self.len - offset,
359            // NOTE: slice key implementation requires us to preserve the
360            // `self.slice` pointer, even if the slice is empty.
361            ptr: unsafe { self.ptr.byte_add(offset) },
362            terminate: T::new(self.terminate.get() && (start <= self.len)),
363            slice: PhantomData,
364        }
365    }
366
367    #[inline]
368    fn common_prefix(self, other: Self) -> Self {
369        let index = r#unsized::common_prefix(self.as_slice(), other.as_slice());
370
371        Self {
372            ptr: self.ptr,
373            len: index,
374            terminate: T::new(
375                self.terminate.get()
376                    && other.terminate.get()
377                    && index == self.len
378                    && index == other.len,
379            ),
380            slice: PhantomData,
381        }
382    }
383}
384
385#[doc(hidden)]
386#[repr(transparent)]
387#[derive(Debug, Default)]
388pub struct Writer(Vec<u8>);
389
390impl Writer {
391    unsafe fn as_slice_unchecked<I: r#unsized::Invariant, R: ?Sized>(&self) -> &Slice<I, R> {
392        let raw = I::Terminate::trim(self.0.as_slice());
393        unsafe { Slice::<I, R>::new_unchecked(core::mem::transmute_copy::<&[u8], &R>(&raw)) }
394    }
395}
396
397impl<'k, T: Terminate> key::Write<Reader<'k, T>> for Writer {
398    type Len = Byte;
399
400    #[inline]
401    fn new(prefix: Reader<'k, T>, key: ribbit::Packed<edge::Le>) -> (Self, Self::Len) {
402        let len = prefix.len() + key.len().into();
403        let mut buffer = Vec::new();
404        buffer.extend_from_slice(prefix.as_slice());
405        if prefix.terminate.get() {
406            buffer.push(u8::MIN);
407            validate_eq!(key.len().bits(), 0);
408        } else {
409            buffer.extend(key);
410        }
411        (Writer(buffer), len)
412    }
413
414    #[inline]
415    fn replace(&mut self, start: Self::Len, node: u8, edge: ribbit::Packed<edge::Le>) -> Self::Len {
416        validate!(start.0 <= self.0.len());
417        self.0.truncate(start.0);
418        self.0.push(node);
419        self.0.extend(edge);
420        Byte(self.0.len())
421    }
422}