Skip to main content

triomphe/
header.rs

1use alloc::alloc::Layout;
2use alloc::boxed::Box;
3use alloc::string::String;
4use alloc::vec::Vec;
5use core::cmp::Ordering;
6use core::iter::{ExactSizeIterator, Iterator};
7use core::marker::PhantomData;
8use core::mem::ManuallyDrop;
9use core::ptr::{self, addr_of_mut};
10
11use crate::{AllocError, OffsetArc};
12
13use super::{Arc, ArcInner};
14
15/// Structure to allow Arc-managing some fixed-sized data and a variably-sized
16/// slice in a single allocation.
17#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash, PartialOrd, Ord)]
18#[repr(C)]
19pub struct HeaderSlice<H, T: ?Sized> {
20    /// The fixed-sized data.
21    pub header: H,
22
23    /// The dynamically-sized data.
24    pub slice: T,
25}
26
27impl<H, T> Arc<HeaderSlice<H, [T]>> {
28    /// Creates an Arc for a HeaderSlice using the given header struct and
29    /// iterator to generate the slice. The resulting Arc will be fat.
30    ///
31    /// **Panics** if the iterator yields a different number of elements than
32    /// reported, or if the iterator itself panicked. In either case, the
33    /// memory is leaked.
34    pub fn from_header_and_iter<I>(header: H, mut items: I) -> Self
35    where
36        I: Iterator<Item = T> + ExactSizeIterator,
37    {
38        let num_items = items.len();
39
40        let inner = Arc::allocate_for_header_and_slice(num_items);
41
42        unsafe {
43            // Write the data.
44            //
45            // Note that any panics here (i.e. from the iterator) are safe, since
46            // we'll just leak the uninitialized memory.
47            ptr::write(addr_of_mut!((*inner.as_ptr()).data.header), header);
48            let mut current = addr_of_mut!((*inner.as_ptr()).data.slice) as *mut T;
49            for _ in 0..num_items {
50                // ZST writes are a no-op, but we still check iterator length
51                ptr::write(
52                    current,
53                    items
54                        .next()
55                        .expect("ExactSizeIterator over-reported length"),
56                );
57                current = current.add(1);
58            }
59            assert!(
60                items.next().is_none(),
61                "ExactSizeIterator under-reported length"
62            );
63        }
64
65        // Safety: ptr is valid & the inner structure is fully initialized
66        Arc {
67            p: inner,
68            phantom: PhantomData,
69        }
70    }
71
72    /// Fallible version of [`Arc::from_header_and_iter`].
73    ///
74    /// Returns `Err(AllocError)` if allocation fails. As with
75    /// [`Arc::from_header_and_iter`], this **panics** (and leaks the
76    /// uninitialized memory) if the iterator yields a different number of
77    /// elements than it reported, or if the iterator itself panics.
78    pub fn try_from_header_and_iter<I>(header: H, mut items: I) -> Result<Self, AllocError>
79    where
80        I: Iterator<Item = T> + ExactSizeIterator,
81    {
82        let num_items = items.len();
83
84        let inner = Arc::try_allocate_for_header_and_slice(num_items)?;
85
86        unsafe {
87            // Write the data.
88            //
89            // Note that any panics here (i.e. from the iterator) are safe, since
90            // we'll just leak the uninitialized memory.
91            ptr::write(addr_of_mut!((*inner.as_ptr()).data.header), header);
92            let mut current = addr_of_mut!((*inner.as_ptr()).data.slice) as *mut T;
93            for _ in 0..num_items {
94                // ZST writes are a no-op, but we still check iterator length
95                ptr::write(
96                    current,
97                    items
98                        .next()
99                        .expect("ExactSizeIterator over-reported length"),
100                );
101                current = current.add(1);
102            }
103            assert!(
104                items.next().is_none(),
105                "ExactSizeIterator under-reported length"
106            );
107        }
108
109        // Safety: ptr is valid & the inner structure is fully initialized
110        Ok(Arc {
111            p: inner,
112            phantom: PhantomData,
113        })
114    }
115
116    /// Creates an Arc for a HeaderSlice using the given header struct and
117    /// slice of copyable items. The items will be copied into the resulting
118    /// Arc, which will be fat.
119    pub fn from_header_and_slice(header: H, items: &[T]) -> Self
120    where
121        T: Copy,
122    {
123        let num_items = items.len();
124
125        let inner = Arc::allocate_for_header_and_slice(num_items);
126
127        unsafe {
128            // Safety
129            // Header is valid (just allocated)
130            ptr::write(addr_of_mut!((*inner.as_ptr()).data.header), header);
131
132            // dst points to `num_items` of uninitialized T's
133            // T: Copy makes bytewise copying safe
134            let dst: *mut [T] = addr_of_mut!((*inner.as_ptr()).data.slice);
135            ptr::copy_nonoverlapping(items.as_ptr(), dst as *mut T, num_items);
136        }
137
138        // Safety: ptr is valid & the inner structure is fully initialized
139        Arc {
140            p: inner,
141            phantom: PhantomData,
142        }
143    }
144
145    /// Fallible version of [`Arc::from_header_and_slice`].
146    ///
147    /// Returns `Err(AllocError)` instead of aborting on allocation failure.
148    pub fn try_from_header_and_slice(header: H, items: &[T]) -> Result<Self, AllocError>
149    where
150        T: Copy,
151    {
152        let num_items = items.len();
153
154        let inner = Arc::try_allocate_for_header_and_slice(num_items)?;
155
156        unsafe {
157            // Safety
158            // Header is valid (just allocated)
159            ptr::write(addr_of_mut!((*inner.as_ptr()).data.header), header);
160
161            // dst points to `num_items` of uninitialized T's
162            // T: Copy makes bytewise copying safe
163            let dst: *mut [T] = addr_of_mut!((*inner.as_ptr()).data.slice);
164            ptr::copy_nonoverlapping(items.as_ptr(), dst as *mut T, num_items);
165        }
166
167        // Safety: ptr is valid & the inner structure is fully initialized
168        Ok(Arc {
169            p: inner,
170            phantom: PhantomData,
171        })
172    }
173
174    /// Creates an Arc for a HeaderSlice using the given header struct and
175    /// vec to generate the slice. The resulting Arc will be fat.
176    pub fn from_header_and_vec(header: H, mut v: Vec<T>) -> Self {
177        let len = v.len();
178
179        let inner = Arc::allocate_for_header_and_slice(len);
180
181        unsafe {
182            // Safety: inner is a valid pointer, so this can't go out of bounds
183            let dst = addr_of_mut!((*inner.as_ptr()).data.header);
184
185            // Safety: `dst` is valid for writes (just allocated)
186            ptr::write(dst, header);
187        }
188
189        unsafe {
190            let src = v.as_mut_ptr();
191
192            // Safety: inner is a valid pointer, so this can't go out of bounds
193            let dst = addr_of_mut!((*inner.as_ptr()).data.slice) as *mut T;
194
195            // Safety:
196            // - `src` is valid for reads for `len` (got from `Vec`)
197            // - `dst` is valid for writes for `len` (just allocated, with layout for appropriate slice)
198            // - `src` and `dst` don't overlap (separate allocations)
199            ptr::copy_nonoverlapping(src, dst, len);
200
201            // Deallocate vec without dropping `T`
202            //
203            // Safety: 0..0 elements are always initialized, 0 <= cap for any cap
204            v.set_len(0);
205        }
206
207        // Safety: ptr is valid & the inner structure is fully initialized
208        Arc {
209            p: inner,
210            phantom: PhantomData,
211        }
212    }
213
214    /// Fallible version of [`Arc::from_header_and_vec`].
215    ///
216    /// Returns `Err(AllocError)` instead of aborting on allocation failure.
217    pub fn try_from_header_and_vec(header: H, mut v: Vec<T>) -> Result<Self, AllocError> {
218        let len = v.len();
219
220        let inner = Arc::try_allocate_for_header_and_slice(len)?;
221
222        unsafe {
223            // Safety: inner is a valid pointer, so this can't go out of bounds
224            let dst = addr_of_mut!((*inner.as_ptr()).data.header);
225
226            // Safety: `dst` is valid for writes (just allocated)
227            ptr::write(dst, header);
228        }
229
230        unsafe {
231            let src = v.as_mut_ptr();
232
233            // Safety: inner is a valid pointer, so this can't go out of bounds
234            let dst = addr_of_mut!((*inner.as_ptr()).data.slice) as *mut T;
235
236            // Safety:
237            // - `src` is valid for reads for `len` (got from `Vec`)
238            // - `dst` is valid for writes for `len` (just allocated, with layout for appropriate slice)
239            // - `src` and `dst` don't overlap (separate allocations)
240            ptr::copy_nonoverlapping(src, dst, len);
241
242            // Deallocate vec without dropping `T`
243            //
244            // Safety: 0..0 elements are always initialized, 0 <= cap for any cap
245            v.set_len(0);
246        }
247
248        // Safety: ptr is valid & the inner structure is fully initialized
249        Ok(Arc {
250            p: inner,
251            phantom: PhantomData,
252        })
253    }
254}
255
256impl<H> Arc<HeaderSlice<H, str>> {
257    /// Creates an Arc for a HeaderSlice using the given header struct and
258    /// a str slice to generate the slice. The resulting Arc will be fat.
259    pub fn from_header_and_str(header: H, string: &str) -> Self {
260        let bytes = Arc::from_header_and_slice(header, string.as_bytes());
261
262        // Safety: `ArcInner` and `HeaderSlice` are `repr(C)`, `str` has the same layout as `[u8]`,
263        //         thus it's ok to "transmute" between `Arc<HeaderSlice<H, [u8]>>` and `Arc<HeaderSlice<H, str>>`.
264        //
265        //         `bytes` are a valid string since we've just got them from a valid `str`.
266        unsafe { Arc::from_raw_inner(Arc::into_raw_inner(bytes) as _) }
267    }
268
269    /// Fallible version of [`Arc::from_header_and_str`].
270    ///
271    /// Returns `Err(AllocError)` instead of aborting on allocation failure.
272    pub fn try_from_header_and_str(header: H, string: &str) -> Result<Self, AllocError> {
273        let bytes = Arc::try_from_header_and_slice(header, string.as_bytes())?;
274
275        // Safety: `ArcInner` and `HeaderSlice` are `repr(C)`, `str` has the same layout as `[u8]`,
276        //         thus it's ok to "transmute" between `Arc<HeaderSlice<H, [u8]>>` and `Arc<HeaderSlice<H, str>>`.
277        //
278        //         `bytes` are a valid string since we've just got them from a valid `str`.
279        Ok(unsafe { Arc::from_raw_inner(Arc::into_raw_inner(bytes) as _) })
280    }
281}
282
283/// Header data with an inline length. Consumers that use HeaderWithLength as the
284/// Header type in HeaderSlice can take advantage of ThinArc.
285#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
286#[repr(C)]
287pub struct HeaderWithLength<H> {
288    /// The fixed-sized data.
289    pub header: H,
290
291    /// The slice length.
292    pub length: usize,
293}
294
295impl<H> HeaderWithLength<H> {
296    /// Creates a new HeaderWithLength.
297    #[inline]
298    pub fn new(header: H, length: usize) -> Self {
299        HeaderWithLength { header, length }
300    }
301}
302
303impl<T: ?Sized> From<Arc<HeaderSlice<(), T>>> for Arc<T> {
304    fn from(this: Arc<HeaderSlice<(), T>>) -> Self {
305        debug_assert_eq!(
306            Layout::for_value::<HeaderSlice<(), T>>(&this),
307            Layout::for_value::<T>(&this.slice)
308        );
309
310        // Safety: `HeaderSlice<(), T>` and `T` has the same layout
311        unsafe { Arc::from_raw_inner(Arc::into_raw_inner(this) as _) }
312    }
313}
314
315impl<T: ?Sized> From<Arc<T>> for Arc<HeaderSlice<(), T>> {
316    fn from(this: Arc<T>) -> Self {
317        // Safety: `T` and `HeaderSlice<(), T>` has the same layout
318        unsafe { Arc::from_raw_inner(Arc::into_raw_inner(this) as _) }
319    }
320}
321
322impl<T: Copy> From<&[T]> for Arc<[T]> {
323    fn from(slice: &[T]) -> Self {
324        Arc::from_header_and_slice((), slice).into()
325    }
326}
327
328impl From<&str> for Arc<str> {
329    fn from(s: &str) -> Self {
330        Arc::from_header_and_str((), s).into()
331    }
332}
333
334impl From<String> for Arc<str> {
335    fn from(s: String) -> Self {
336        Self::from(&s[..])
337    }
338}
339
340impl From<&str> for OffsetArc<str> {
341    fn from(s: &str) -> Self {
342        Arc::into_raw_offset(Arc::from(s))
343    }
344}
345
346impl From<String> for OffsetArc<str> {
347    fn from(s: String) -> Self {
348        Self::from(&s[..])
349    }
350}
351
352// FIXME: once `pointer::with_metadata_of` is stable or
353//        implementable on stable without assuming ptr layout
354//        this will be able to accept `T: ?Sized`.
355impl<T> From<Box<T>> for Arc<T> {
356    fn from(b: Box<T>) -> Self {
357        let layout = Layout::for_value::<T>(&b);
358
359        // Safety: the closure only changes the type of the pointer
360        let inner = unsafe { Self::allocate_for_layout(layout, |mem| mem as *mut ArcInner<T>) };
361
362        unsafe {
363            let src = Box::into_raw(b);
364
365            // Safety: inner is a valid pointer, so this can't go out of bounds
366            let dst = addr_of_mut!((*inner.as_ptr()).data);
367
368            // Safety:
369            // - `src` is valid for reads (got from `Box`)
370            // - `dst` is valid for writes (just allocated)
371            // - `src` and `dst` don't overlap (separate allocations)
372            ptr::copy_nonoverlapping(src, dst, 1);
373
374            // Deallocate box without dropping `T`
375            //
376            // Safety:
377            // - `src` has been got from `Box::into_raw`
378            // - `ManuallyDrop<T>` is guaranteed to have the same layout as `T`
379            drop(Box::<ManuallyDrop<T>>::from_raw(src as _));
380        }
381
382        Arc {
383            p: inner,
384            phantom: PhantomData,
385        }
386    }
387}
388
389impl<T> From<Vec<T>> for Arc<[T]> {
390    fn from(v: Vec<T>) -> Self {
391        Arc::from_header_and_vec((), v).into()
392    }
393}
394
395/// A type wrapping `HeaderSlice<HeaderWithLength<H>, T>` that is used internally in `ThinArc`.
396///
397/// # Safety
398///
399/// Safety-usable invariants:
400///
401/// - This is guaranteed to have the same representation as `HeaderSlice<HeaderWithLength<H>, [T]>`
402/// - The header length (`.length()`) is checked to be the slice length
403#[derive(Debug, Hash, Eq, PartialEq, Ord, PartialOrd)]
404#[repr(transparent)]
405pub struct HeaderSliceWithLengthProtected<H, T> {
406    // Invariant: the header's length field must be the slice length
407    inner: HeaderSliceWithLengthUnchecked<H, T>,
408}
409
410pub(crate) type HeaderSliceWithLengthUnchecked<H, T> = HeaderSlice<HeaderWithLength<H>, [T]>;
411
412impl<H, T> HeaderSliceWithLengthProtected<H, T> {
413    pub fn header(&self) -> &H {
414        &self.inner.header.header
415    }
416    pub fn header_mut(&mut self) -> &mut H {
417        // Safety: only the length is unsafe to mutate
418        &mut self.inner.header.header
419    }
420    pub fn length(&self) -> usize {
421        self.inner.header.length
422    }
423
424    pub fn slice(&self) -> &[T] {
425        &self.inner.slice
426    }
427    pub fn slice_mut(&mut self) -> &mut [T] {
428        // Safety: only the length is unsafe to mutate
429        &mut self.inner.slice
430    }
431    pub(crate) fn inner(&self) -> &HeaderSliceWithLengthUnchecked<H, T> {
432        // This is safe in an immutable context
433        &self.inner
434    }
435}
436
437impl<H: PartialOrd, T: ?Sized + PartialOrd> PartialOrd for HeaderSlice<HeaderWithLength<H>, T> {
438    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
439        (&self.header.header, &self.slice).partial_cmp(&(&other.header.header, &other.slice))
440    }
441}
442
443impl<H: Ord, T: ?Sized + Ord> Ord for HeaderSlice<HeaderWithLength<H>, T> {
444    fn cmp(&self, other: &Self) -> Ordering {
445        (&self.header.header, &self.slice).cmp(&(&other.header.header, &other.slice))
446    }
447}
448
449#[cfg(test)]
450mod tests {
451    use alloc::boxed::Box;
452    use alloc::string::String;
453    use alloc::vec;
454    use core::iter;
455
456    use crate::{Arc, HeaderSlice};
457
458    #[test]
459    fn from_header_and_iter_smoke() {
460        let arc = Arc::from_header_and_iter(
461            (42u32, 17u8),
462            IntoIterator::into_iter([1u16, 2, 3, 4, 5, 6, 7]),
463        );
464
465        assert_eq!(arc.header, (42, 17));
466        assert_eq!(arc.slice, [1, 2, 3, 4, 5, 6, 7]);
467    }
468
469    #[test]
470    fn try_from_header_and_iter_smoke() {
471        let arc = Arc::try_from_header_and_iter(
472            (42u32, 17u8),
473            IntoIterator::into_iter([1u16, 2, 3, 4, 5, 6, 7]),
474        )
475        .unwrap();
476
477        assert_eq!(arc.header, (42, 17));
478        assert_eq!(arc.slice, [1, 2, 3, 4, 5, 6, 7]);
479    }
480
481    #[test]
482    fn from_header_and_slice_smoke() {
483        let arc = Arc::from_header_and_slice((42u32, 17u8), &[1u16, 2, 3, 4, 5, 6, 7]);
484
485        assert_eq!(arc.header, (42, 17));
486        assert_eq!(arc.slice, [1u16, 2, 3, 4, 5, 6, 7]);
487    }
488
489    #[test]
490    fn try_from_header_and_slice_smoke() {
491        let arc = Arc::try_from_header_and_slice((42u32, 17u8), &[1u16, 2, 3, 4, 5, 6, 7]).unwrap();
492
493        assert_eq!(arc.header, (42, 17));
494        assert_eq!(arc.slice, [1u16, 2, 3, 4, 5, 6, 7]);
495    }
496
497    #[test]
498    fn try_from_header_and_vec_smoke() {
499        let arc =
500            Arc::try_from_header_and_vec((42u32, 17u8), vec![1u16, 2, 3, 4, 5, 6, 7]).unwrap();
501
502        assert_eq!(arc.header, (42, 17));
503        assert_eq!(arc.slice, [1u16, 2, 3, 4, 5, 6, 7]);
504    }
505
506    #[test]
507    fn try_from_header_and_str_smoke() {
508        let a = Arc::try_from_header_and_str(
509            42,
510            "The answer to the ultimate question of life, the universe, and everything",
511        )
512        .unwrap();
513        assert_eq!(a.header, 42);
514        assert_eq!(
515            &a.slice,
516            "The answer to the ultimate question of life, the universe, and everything"
517        );
518
519        let empty = Arc::try_from_header_and_str((), "").unwrap();
520        assert_eq!(&empty.slice, "");
521    }
522
523    #[test]
524    fn from_header_and_vec_smoke() {
525        let arc = Arc::from_header_and_vec((42u32, 17u8), vec![1u16, 2, 3, 4, 5, 6, 7]);
526
527        assert_eq!(arc.header, (42, 17));
528        assert_eq!(arc.slice, [1u16, 2, 3, 4, 5, 6, 7]);
529    }
530
531    #[test]
532    fn from_header_and_iter_empty() {
533        let arc = Arc::from_header_and_iter((42u32, 17u8), iter::empty::<u16>());
534
535        assert_eq!(arc.header, (42, 17));
536        assert_eq!(arc.slice, []);
537    }
538
539    #[test]
540    fn from_header_and_slice_empty() {
541        let arc = Arc::from_header_and_slice((42u32, 17u8), &[1u16; 0]);
542
543        assert_eq!(arc.header, (42, 17));
544        assert_eq!(arc.slice, []);
545    }
546
547    #[test]
548    fn from_header_and_vec_empty() {
549        let arc = Arc::from_header_and_vec((42u32, 17u8), vec![1u16; 0]);
550
551        assert_eq!(arc.header, (42, 17));
552        assert_eq!(arc.slice, []);
553    }
554
555    #[test]
556    fn issue_13_empty() {
557        crate::Arc::from_header_and_iter((), iter::empty::<usize>());
558    }
559
560    #[test]
561    fn issue_13_consumption() {
562        let s: &[u8] = &[0u8; 255];
563        crate::Arc::from_header_and_iter((), s.iter().copied());
564    }
565
566    #[test]
567    fn from_header_and_str_smoke() {
568        let a = Arc::from_header_and_str(
569            42,
570            "The answer to the ultimate question of life, the universe, and everything",
571        );
572        assert_eq!(a.header, 42);
573        assert_eq!(
574            &a.slice,
575            "The answer to the ultimate question of life, the universe, and everything"
576        );
577
578        let empty = Arc::from_header_and_str((), "");
579        assert_eq!(&empty.slice, "");
580    }
581
582    #[test]
583    fn erase_and_create_from_thin_air_header() {
584        let a: Arc<HeaderSlice<(), [u32]>> = Arc::from_header_and_slice((), &[12, 17, 16]);
585        let b: Arc<[u32]> = a.into();
586
587        assert_eq!(&*b, [12, 17, 16]);
588
589        let c: Arc<HeaderSlice<(), [u32]>> = b.into();
590
591        assert_eq!(&c.slice, [12, 17, 16]);
592    }
593
594    #[test]
595    fn from_box_and_vec() {
596        let b = Box::new(String::from("xxx"));
597        let b = Arc::<String>::from(b);
598        assert_eq!(&*b, "xxx");
599
600        let v = vec![String::from("1"), String::from("2"), String::from("3")];
601        let v = Arc::<[_]>::from(v);
602        assert_eq!(
603            &*v,
604            [String::from("1"), String::from("2"), String::from("3")]
605        );
606
607        let mut v = vec![String::from("1"), String::from("2"), String::from("3")];
608        v.reserve(10);
609        let v = Arc::<[_]>::from(v);
610        assert_eq!(
611            &*v,
612            [String::from("1"), String::from("2"), String::from("3")]
613        );
614    }
615
616    /// It’s possible to make a generic `Arc` wrapper that supports both:
617    ///
618    /// * `T: !Sized`
619    /// * `Arc::make_mut` if `T: Sized`
620    #[test]
621    fn dst_and_make_mut() {
622        struct MyArc<T: ?Sized>(Arc<HeaderSlice<MyHeader, T>>);
623
624        #[derive(Clone)]
625        struct MyHeader {
626            // Very interesting things go here
627        }
628
629        // MyArc<str> is possible
630        let dst: MyArc<str> = MyArc(Arc::from_header_and_str(MyHeader {}, "example"));
631        assert_eq!(&dst.0.slice, "example");
632
633        // `make_mut` is still available when `T: Sized`
634        let mut answer: MyArc<u32> = MyArc(Arc::new(HeaderSlice {
635            header: MyHeader {},
636            // Not actually a slice in this case,
637            // but `HeaderSlice` is required to use `from_header_and_str`
638            // and we want the same `MyArc` to support both cases.
639            slice: 6 * 9,
640        }));
641        let mut_ref = Arc::make_mut(&mut answer.0);
642        mut_ref.slice = 42;
643    }
644}