1 #![cfg(feature = "extern_crate_alloc")]
2
3 //! Stuff to boost things in the `alloc` crate.
4 //!
5 //! * You must enable the `extern_crate_alloc` feature of `bytemuck` or you will
6 //! not be able to use this module!
7
8 use super::*;
9 use alloc::{
10 alloc::{alloc_zeroed, Layout},
11 boxed::Box,
12 vec,
13 vec::Vec,
14 };
15 use core::convert::TryInto;
16
17 /// As [`try_cast_box`](try_cast_box), but unwraps for you.
18 #[inline]
cast_box<A: Pod, B: Pod>(input: Box<A>) -> Box<B>19 pub fn cast_box<A: Pod, B: Pod>(input: Box<A>) -> Box<B> {
20 try_cast_box(input).map_err(|(e, _v)| e).unwrap()
21 }
22
23 /// Attempts to cast the content type of a [`Box`](alloc::boxed::Box).
24 ///
25 /// On failure you get back an error along with the starting `Box`.
26 ///
27 /// ## Failure
28 ///
29 /// * The start and end content type of the `Box` must have the exact same
30 /// alignment.
31 /// * The start and end size of the `Box` must have the exact same size.
32 #[inline]
try_cast_box<A: Pod, B: Pod>( input: Box<A>, ) -> Result<Box<B>, (PodCastError, Box<A>)>33 pub fn try_cast_box<A: Pod, B: Pod>(
34 input: Box<A>,
35 ) -> Result<Box<B>, (PodCastError, Box<A>)> {
36 if align_of::<A>() != align_of::<B>() {
37 Err((PodCastError::AlignmentMismatch, input))
38 } else if size_of::<A>() != size_of::<B>() {
39 Err((PodCastError::SizeMismatch, input))
40 } else {
41 // Note(Lokathor): This is much simpler than with the Vec casting!
42 let ptr: *mut B = Box::into_raw(input) as *mut B;
43 Ok(unsafe { Box::from_raw(ptr) })
44 }
45 }
46
47 /// Allocates a `Box<T>` with all of the contents being zeroed out.
48 ///
49 /// This uses the global allocator to create a zeroed allocation and _then_
50 /// turns it into a Box. In other words, it's 100% assured that the zeroed data
51 /// won't be put temporarily on the stack. You can make a box of any size
52 /// without fear of a stack overflow.
53 ///
54 /// ## Failure
55 ///
56 /// This fails if the allocation fails.
57 #[inline]
try_zeroed_box<T: Zeroable>() -> Result<Box<T>, ()>58 pub fn try_zeroed_box<T: Zeroable>() -> Result<Box<T>, ()> {
59 if size_of::<T>() == 0 {
60 // This will not allocate but simple create a dangling slice pointer.
61 // NB: We go the way via a push to `Vec<T>` to ensure the compiler
62 // does not allocate space for T on the stack even if the branch
63 // would not be taken.
64 let mut vec = Vec::with_capacity(1);
65 vec.resize_with(1, || T::zeroed());
66 let ptr: Box<[T; 1]> = vec.into_boxed_slice().try_into().ok().unwrap();
67 debug_assert!(
68 align_of::<[T; 1]>() == align_of::<T>()
69 && size_of::<[T; 1]>() == size_of::<T>()
70 );
71 // NB: We basically do the same as in try_cast_box here:
72 let ptr: Box<T> = unsafe { Box::from_raw(Box::into_raw(ptr) as *mut _) };
73 return Ok(ptr);
74 }
75 let layout =
76 Layout::from_size_align(size_of::<T>(), align_of::<T>()).unwrap();
77 let ptr = unsafe { alloc_zeroed(layout) };
78 if ptr.is_null() {
79 // we don't know what the error is because `alloc_zeroed` is a dumb API
80 Err(())
81 } else {
82 Ok(unsafe { Box::<T>::from_raw(ptr as *mut T) })
83 }
84 }
85
86 /// As [`try_zeroed_box`], but unwraps for you.
87 #[inline]
zeroed_box<T: Zeroable>() -> Box<T>88 pub fn zeroed_box<T: Zeroable>() -> Box<T> {
89 try_zeroed_box().unwrap()
90 }
91
92 /// Allocates a `Box<[T]>` with all contents being zeroed out.
93 ///
94 /// This uses the global allocator to create a zeroed allocation and _then_
95 /// turns it into a Box. In other words, it's 100% assured that the zeroed data
96 /// won't be put temporarily on the stack. You can make a box of any size
97 /// without fear of a stack overflow.
98 ///
99 /// ## Failure
100 ///
101 /// This fails if the allocation fails.
102 #[inline]
try_zeroed_slice_box<T: Zeroable>( length: usize, ) -> Result<Box<[T]>, ()>103 pub fn try_zeroed_slice_box<T: Zeroable>(
104 length: usize,
105 ) -> Result<Box<[T]>, ()> {
106 if size_of::<T>() == 0 {
107 // This will not allocate but simple create a dangling slice pointer.
108 let mut vec = Vec::with_capacity(length);
109 vec.resize_with(length, || T::zeroed());
110 return Ok(vec.into_boxed_slice());
111 }
112 if length == 0 {
113 // This will also not allocate.
114 return Ok(Vec::new().into_boxed_slice());
115 }
116 // For Pod types, the layout of the array/slice is equivalent to repeating the
117 // type.
118 let layout_length = size_of::<T>().checked_mul(length).ok_or(())?;
119 assert!(layout_length != 0);
120 let layout =
121 Layout::from_size_align(layout_length, align_of::<T>()).map_err(|_| ())?;
122 let ptr = unsafe { alloc_zeroed(layout) };
123 if ptr.is_null() {
124 // we don't know what the error is because `alloc_zeroed` is a dumb API
125 Err(())
126 } else {
127 let slice =
128 unsafe { core::slice::from_raw_parts_mut(ptr as *mut T, length) };
129 Ok(unsafe { Box::<[T]>::from_raw(slice) })
130 }
131 }
132
133 /// As [`try_zeroed_slice_box`](try_zeroed_slice_box), but unwraps for you.
zeroed_slice_box<T: Zeroable>(length: usize) -> Box<[T]>134 pub fn zeroed_slice_box<T: Zeroable>(length: usize) -> Box<[T]> {
135 try_zeroed_slice_box(length).unwrap()
136 }
137
138 /// As [`try_cast_vec`](try_cast_vec), but unwraps for you.
139 #[inline]
cast_vec<A: Pod, B: Pod>(input: Vec<A>) -> Vec<B>140 pub fn cast_vec<A: Pod, B: Pod>(input: Vec<A>) -> Vec<B> {
141 try_cast_vec(input).map_err(|(e, _v)| e).unwrap()
142 }
143
144 /// Attempts to cast the content type of a [`Vec`](alloc::vec::Vec).
145 ///
146 /// On failure you get back an error along with the starting `Vec`.
147 ///
148 /// ## Failure
149 ///
150 /// * The start and end content type of the `Vec` must have the exact same
151 /// alignment.
152 /// * The start and end size of the `Vec` must have the exact same size.
153 /// * In the future this second restriction might be lessened by having the
154 /// capacity and length get adjusted during transmutation, but for now it's
155 /// absolute.
156 #[inline]
try_cast_vec<A: Pod, B: Pod>( input: Vec<A>, ) -> Result<Vec<B>, (PodCastError, Vec<A>)>157 pub fn try_cast_vec<A: Pod, B: Pod>(
158 input: Vec<A>,
159 ) -> Result<Vec<B>, (PodCastError, Vec<A>)> {
160 if align_of::<A>() != align_of::<B>() {
161 Err((PodCastError::AlignmentMismatch, input))
162 } else if size_of::<A>() != size_of::<B>() {
163 // Note(Lokathor): Under some conditions it would be possible to cast
164 // between Vec content types of the same alignment but different sizes by
165 // changing the capacity and len values in the output Vec. However, we will
166 // not attempt that for now.
167 Err((PodCastError::SizeMismatch, input))
168 } else {
169 // Note(Lokathor): First we record the length and capacity, which don't have
170 // any secret provenance metadata.
171 let length: usize = input.len();
172 let capacity: usize = input.capacity();
173 // Note(Lokathor): Next we "pre-forget" the old Vec by wrapping with
174 // ManuallyDrop, because if we used `core::mem::forget` after taking the
175 // pointer then that would invalidate our pointer. In nightly there's a
176 // "into raw parts" method, which we can switch this too eventually.
177 let mut manual_drop_vec = ManuallyDrop::new(input);
178 let vec_ptr: *mut A = manual_drop_vec.as_mut_ptr();
179 let ptr: *mut B = vec_ptr as *mut B;
180 Ok(unsafe { Vec::from_raw_parts(ptr, length, capacity) })
181 }
182 }
183
184 /// This "collects" a slice of pod data into a vec of a different pod type.
185 ///
186 /// Unlike with [`cast_slice`] and [`cast_slice_mut`], this will always work.
187 ///
188 /// The output vec will be of a minimal size/capacity to hold the slice given.
189 ///
190 /// ```rust
191 /// # use bytemuck::*;
192 /// let halfwords: [u16; 4] = [5, 6, 7, 8];
193 /// let vec_of_words: Vec<u32> = pod_collect_to_vec(&halfwords);
194 /// if cfg!(target_endian = "little") {
195 /// assert_eq!(&vec_of_words[..], &[0x0006_0005, 0x0008_0007][..])
196 /// } else {
197 /// assert_eq!(&vec_of_words[..], &[0x0005_0006, 0x0007_0008][..])
198 /// }
199 /// ```
pod_collect_to_vec<A: Pod, B: Pod>(src: &[A]) -> Vec<B>200 pub fn pod_collect_to_vec<A: Pod, B: Pod>(src: &[A]) -> Vec<B> {
201 let src_size = size_of_val(src);
202 // Note(Lokathor): dst_count is rounded up so that the dest will always be at
203 // least as many bytes as the src.
204 let dst_count = src_size / size_of::<B>()
205 + if src_size % size_of::<B>() != 0 { 1 } else { 0 };
206 let mut dst = vec![B::zeroed(); dst_count];
207
208 let src_bytes: &[u8] = cast_slice(src);
209 let dst_bytes: &mut [u8] = cast_slice_mut(&mut dst[..]);
210 dst_bytes[..src_size].copy_from_slice(src_bytes);
211 dst
212 }
213