1 // Copyright 2019 The ChromiumOS Authors 2 // Use of this source code is governed by a BSD-style license that can be 3 // found in the LICENSE file. 4 5 use std::alloc::alloc; 6 use std::alloc::alloc_zeroed; 7 use std::alloc::dealloc; 8 use std::alloc::Layout; 9 use std::cmp::min; 10 11 /// A contiguous memory allocation with a specified size and alignment, with a 12 /// Drop impl to perform the deallocation. 13 /// 14 /// Conceptually this is like a Box<[u8]> but for which we can select a minimum 15 /// required alignment at the time of allocation. 16 /// 17 /// # Example 18 /// 19 /// ``` 20 /// use std::alloc::Layout; 21 /// use std::mem; 22 /// use base::LayoutAllocation; 23 /// 24 /// #[repr(C)] 25 /// struct Header { 26 /// q: usize, 27 /// entries: [Entry; 0], // flexible array member 28 /// } 29 /// 30 /// #[repr(C)] 31 /// struct Entry { 32 /// e: usize, 33 /// } 34 /// 35 /// fn demo(num_entries: usize) { 36 /// let size = mem::size_of::<Header>() + num_entries * mem::size_of::<Entry>(); 37 /// let layout = Layout::from_size_align(size, mem::align_of::<Header>()).unwrap(); 38 /// let mut allocation = LayoutAllocation::zeroed(layout); 39 /// 40 /// // SAFETY: 41 /// // Safe to obtain an exclusive reference because there are no other 42 /// // references to the allocation yet and all-zero is a valid bit pattern for 43 /// // our header. 44 /// let header = unsafe { allocation.as_mut::<Header>() }; 45 /// } 46 /// ``` 47 pub struct LayoutAllocation { 48 ptr: *mut u8, 49 layout: Layout, 50 } 51 52 impl LayoutAllocation { 53 /// Allocates memory with the specified size and alignment. The content is 54 /// not initialized. 55 /// 56 /// Uninitialized data is not safe to read. Further, it is not safe to 57 /// obtain a reference to data potentially holding a bit pattern 58 /// incompatible with its type, for example an uninitialized bool or enum. uninitialized(layout: Layout) -> Self59 pub fn uninitialized(layout: Layout) -> Self { 60 let ptr = if layout.size() > 0 { 61 // SAFETY: 62 // Safe as long as we guarantee layout.size() > 0. 63 unsafe { alloc(layout) } 64 } else { 65 layout.align() as *mut u8 66 }; 67 LayoutAllocation { ptr, layout } 68 } 69 70 /// Allocates memory with the specified size and alignment and initializes 71 /// the content to all zero-bytes. 72 /// 73 /// Note that zeroing the memory does not necessarily make it safe to obtain 74 /// a reference to the allocation. Depending on the intended type T, 75 /// all-zero may or may not be a legal bit pattern for that type. For 76 /// example obtaining a reference would immediately be undefined behavior if 77 /// one of the fields has type NonZeroUsize. zeroed(layout: Layout) -> Self78 pub fn zeroed(layout: Layout) -> Self { 79 let ptr = if layout.size() > 0 { 80 // SAFETY: 81 // Safe as long as we guarantee layout.size() > 0. 82 unsafe { alloc_zeroed(layout) } 83 } else { 84 layout.align() as *mut u8 85 }; 86 LayoutAllocation { ptr, layout } 87 } 88 89 /// Returns a raw pointer to the allocated data. as_ptr<T>(&self) -> *mut T90 pub fn as_ptr<T>(&self) -> *mut T { 91 self.ptr as *mut T 92 } 93 94 /// Returns a reference to the `Layout` used to create this allocation. layout(&self) -> &Layout95 pub fn layout(&self) -> &Layout { 96 &self.layout 97 } 98 99 /// Returns a shared reference to the allocated data. 100 /// 101 /// # Safety 102 /// 103 /// Caller is responsible for ensuring that the data behind this pointer has 104 /// been initialized as much as necessary and that there are no already 105 /// existing mutable references to any part of the data. as_ref<T>(&self) -> &T106 pub unsafe fn as_ref<T>(&self) -> &T { 107 &*self.as_ptr() 108 } 109 110 /// Returns an exclusive reference to the allocated data. 111 /// 112 /// # Safety 113 /// 114 /// Caller is responsible for ensuring that the data behind this pointer has 115 /// been initialized as much as necessary and that there are no already 116 /// existing references to any part of the data. as_mut<T>(&mut self) -> &mut T117 pub unsafe fn as_mut<T>(&mut self) -> &mut T { 118 &mut *self.as_ptr() 119 } 120 121 /// Returns a shared slice reference to the allocated data. 122 /// 123 /// # Arguments 124 /// 125 /// `num_elements` - Number of `T` elements to include in the slice. 126 /// The length of the slice will be capped to the allocation's size. 127 /// Caller must ensure that any sliced elements are initialized. 128 /// 129 /// # Safety 130 /// 131 /// Caller is responsible for ensuring that the data behind this pointer has 132 /// been initialized as much as necessary and that there are no already 133 /// existing mutable references to any part of the data. as_slice<T>(&self, num_elements: usize) -> &[T]134 pub unsafe fn as_slice<T>(&self, num_elements: usize) -> &[T] { 135 let len = min(num_elements, self.layout.size() / std::mem::size_of::<T>()); 136 std::slice::from_raw_parts(self.as_ptr(), len) 137 } 138 139 /// Returns an exclusive slice reference to the allocated data. 140 /// 141 /// # Arguments 142 /// 143 /// `num_elements` - Number of `T` elements to include in the slice. 144 /// The length of the slice will be capped to the allocation's size. 145 /// Caller must ensure that any sliced elements are initialized. 146 /// 147 /// # Safety 148 /// 149 /// Caller is responsible for ensuring that the data behind this pointer has 150 /// been initialized as much as necessary and that there are no already 151 /// existing references to any part of the data. as_mut_slice<T>(&mut self, num_elements: usize) -> &mut [T]152 pub unsafe fn as_mut_slice<T>(&mut self, num_elements: usize) -> &mut [T] { 153 let len = min(num_elements, self.layout.size() / std::mem::size_of::<T>()); 154 std::slice::from_raw_parts_mut(self.as_ptr(), len) 155 } 156 } 157 158 impl Drop for LayoutAllocation { drop(&mut self)159 fn drop(&mut self) { 160 if self.layout.size() > 0 { 161 // SAFETY: 162 // Safe as long as we guarantee layout.size() > 0. 163 unsafe { 164 dealloc(self.ptr, self.layout); 165 } 166 } 167 } 168 } 169 170 #[cfg(test)] 171 mod tests { 172 use std::mem::align_of; 173 use std::mem::size_of; 174 175 use super::*; 176 177 #[test] test_as_slice_u32()178 fn test_as_slice_u32() { 179 let layout = Layout::from_size_align(size_of::<u32>() * 15, align_of::<u32>()).unwrap(); 180 let allocation = LayoutAllocation::zeroed(layout); 181 // SAFETY: 182 // Slice less than the allocation size, which will return a slice of only the requested 183 // length. 184 let slice: &[u32] = unsafe { allocation.as_slice(15) }; 185 assert_eq!(slice.len(), 15); 186 assert_eq!(slice[0], 0); 187 assert_eq!(slice[14], 0); 188 } 189 190 #[test] test_as_slice_u32_smaller_len()191 fn test_as_slice_u32_smaller_len() { 192 let layout = Layout::from_size_align(size_of::<u32>() * 15, align_of::<u32>()).unwrap(); 193 let allocation = LayoutAllocation::zeroed(layout); 194 195 // SAFETY: 196 // Slice less than the allocation size, which will return a slice of only the requested 197 // length. 198 let slice: &[u32] = unsafe { allocation.as_slice(5) }; 199 assert_eq!(slice.len(), 5); 200 } 201 202 #[test] test_as_slice_u32_larger_len()203 fn test_as_slice_u32_larger_len() { 204 let layout = Layout::from_size_align(size_of::<u32>() * 15, align_of::<u32>()).unwrap(); 205 let allocation = LayoutAllocation::zeroed(layout); 206 207 // SAFETY: 208 // Slice more than the allocation size, which will clamp the returned slice len to the 209 // limit. 210 let slice: &[u32] = unsafe { allocation.as_slice(100) }; 211 assert_eq!(slice.len(), 15); 212 } 213 214 #[test] test_as_slice_u32_remainder()215 fn test_as_slice_u32_remainder() { 216 // Allocate a buffer that is not a multiple of u32 in size. 217 let layout = Layout::from_size_align(size_of::<u32>() * 15 + 2, align_of::<u32>()).unwrap(); 218 let allocation = LayoutAllocation::zeroed(layout); 219 220 // SAFETY: 221 // Slice as many u32s as possible, which should return a slice that only includes the full 222 // u32s, not the trailing 2 bytes. 223 let slice: &[u32] = unsafe { allocation.as_slice(100) }; 224 assert_eq!(slice.len(), 15); 225 } 226 } 227