|
| 1 | +/* |
| 2 | + * Copyright 2018 The Starlark in Rust Authors. |
| 3 | + * Copyright (c) Facebook, Inc. and its affiliates. |
| 4 | + * |
| 5 | + * Licensed under the Apache License, Version 2.0 (the "License"); |
| 6 | + * you may not use this file except in compliance with the License. |
| 7 | + * You may obtain a copy of the License at |
| 8 | + * |
| 9 | + * https://www.apache.org/licenses/LICENSE-2.0 |
| 10 | + * |
| 11 | + * Unless required by applicable law or agreed to in writing, software |
| 12 | + * distributed under the License is distributed on an "AS IS" BASIS, |
| 13 | + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 14 | + * See the License for the specific language governing permissions and |
| 15 | + * limitations under the License. |
| 16 | + */ |
| 17 | + |
| 18 | +//! This type is a copy-paste of `buck2_util::thin_box::ThinBoxSlice`, with some mild adjustments. |
| 19 | +//! |
| 20 | +//! Specifically: |
| 21 | +//! 1. This type guarantees that it's always a pointer with the bottom bit zero. |
| 22 | +//! 2. This type is not implicitly dropped - `run_drop` must be called explicitly. |
| 23 | +
|
| 24 | +use std::alloc; |
| 25 | +use std::alloc::Layout; |
| 26 | +use std::fmt::Debug; |
| 27 | +use std::hash::Hash; |
| 28 | +use std::hash::Hasher; |
| 29 | +use std::mem; |
| 30 | +use std::mem::MaybeUninit; |
| 31 | +use std::ops::Deref; |
| 32 | +use std::ops::DerefMut; |
| 33 | +use std::ptr; |
| 34 | +use std::ptr::NonNull; |
| 35 | +use std::slice; |
| 36 | + |
| 37 | +use allocative::Allocative; |
| 38 | + |
| 39 | +#[repr(C)] |
| 40 | +struct ThinBoxSliceLayout<T> { |
| 41 | + len: usize, |
| 42 | + data: [T; 0], |
| 43 | +} |
| 44 | + |
| 45 | +impl<T> ThinBoxSliceLayout<T> { |
| 46 | + fn offset_of_data() -> usize { |
| 47 | + mem::offset_of!(ThinBoxSliceLayout::<T>, data) |
| 48 | + } |
| 49 | +} |
| 50 | + |
| 51 | +/// `Box<[T]>` but thin pointer. |
| 52 | +/// |
| 53 | +/// Statically allocated for empty slice. |
| 54 | +// We don't really need `'static` here, but we hit type checker limitations. |
| 55 | +pub(super) struct AllocatedThinBoxSlice<T: 'static> { |
| 56 | + /// Pointer to the first element, `ThinBoxSliceLayout.data`. |
| 57 | + ptr: NonNull<T>, |
| 58 | +} |
| 59 | + |
| 60 | +unsafe impl<T: Sync> Sync for AllocatedThinBoxSlice<T> {} |
| 61 | +unsafe impl<T: Send> Send for AllocatedThinBoxSlice<T> {} |
| 62 | + |
| 63 | +impl<T: 'static> AllocatedThinBoxSlice<T> { |
| 64 | + #[inline] |
| 65 | + pub(super) const fn empty() -> AllocatedThinBoxSlice<T> { |
| 66 | + const fn instance<T>() -> &'static ThinBoxSliceLayout<T> { |
| 67 | + assert!(mem::size_of::<ThinBoxSliceLayout<T>>() <= mem::size_of::<u128>()); |
| 68 | + assert!(mem::align_of::<ThinBoxSliceLayout<T>>() <= mem::align_of::<u128>()); |
| 69 | + // Instead of just statically allocating a `ThinBoxSliceLayout<T>` we allocate a |
| 70 | + // `0_u128`. The reason for this is a rustc bug around const UB checks that otherwise |
| 71 | + // incorrectly fires: https://github.com/rust-lang/rust/issues/133523 |
| 72 | + // |
| 73 | + // SAFETY: We just checked that the layout is small enough to fit in a u128. |
| 74 | + unsafe { &*ptr::from_ref(&0u128).cast::<ThinBoxSliceLayout<T>>() } |
| 75 | + } |
| 76 | + |
| 77 | + unsafe { |
| 78 | + AllocatedThinBoxSlice { |
| 79 | + ptr: NonNull::new_unchecked(instance::<T>().data.as_ptr() as *mut T), |
| 80 | + } |
| 81 | + } |
| 82 | + } |
| 83 | + |
| 84 | + /// Allocation layout for a slice of length `len`. |
| 85 | + #[inline] |
| 86 | + fn layout_for_len(len: usize) -> Layout { |
| 87 | + let (layout, _offset_of_data) = Layout::new::<ThinBoxSliceLayout<T>>() |
| 88 | + .extend(Layout::array::<T>(len).unwrap()) |
| 89 | + .unwrap(); |
| 90 | + layout |
| 91 | + } |
| 92 | + |
| 93 | + /// Length of the slice. |
| 94 | + // Not called `len` to avoid overload with `Deref::len`. |
| 95 | + #[inline] |
| 96 | + fn read_len(&self) -> usize { |
| 97 | + unsafe { |
| 98 | + (*self |
| 99 | + .ptr |
| 100 | + .as_ptr() |
| 101 | + .cast::<u8>() |
| 102 | + .sub(ThinBoxSliceLayout::<T>::offset_of_data()) |
| 103 | + .cast::<ThinBoxSliceLayout<T>>()) |
| 104 | + .len |
| 105 | + } |
| 106 | + } |
| 107 | + |
| 108 | + /// Allocate uninitialized memory for a slice of length `len`. |
| 109 | + #[inline] |
| 110 | + pub(super) fn new_uninit(len: usize) -> AllocatedThinBoxSlice<MaybeUninit<T>> { |
| 111 | + if len == 0 { |
| 112 | + AllocatedThinBoxSlice::empty() |
| 113 | + } else { |
| 114 | + let layout = Self::layout_for_len(len); |
| 115 | + unsafe { |
| 116 | + let alloc = alloc::alloc(layout); |
| 117 | + if alloc.is_null() { |
| 118 | + alloc::handle_alloc_error(layout); |
| 119 | + } |
| 120 | + ptr::write(alloc as *mut usize, len); |
| 121 | + let ptr = alloc.add(mem::size_of::<usize>()) as *mut MaybeUninit<T>; |
| 122 | + let ptr = NonNull::new_unchecked(ptr); |
| 123 | + AllocatedThinBoxSlice { ptr } |
| 124 | + } |
| 125 | + } |
| 126 | + } |
| 127 | + |
| 128 | + pub const unsafe fn into_inner(self) -> NonNull<T> { |
| 129 | + self.ptr |
| 130 | + } |
| 131 | + |
| 132 | + pub const unsafe fn from_inner(ptr: NonNull<T>) -> Self { |
| 133 | + Self { ptr } |
| 134 | + } |
| 135 | +} |
| 136 | + |
| 137 | +impl<T: 'static> Deref for AllocatedThinBoxSlice<T> { |
| 138 | + type Target = [T]; |
| 139 | + |
| 140 | + #[inline] |
| 141 | + fn deref(&self) -> &Self::Target { |
| 142 | + unsafe { slice::from_raw_parts(self.ptr.as_ptr(), self.read_len()) } |
| 143 | + } |
| 144 | +} |
| 145 | + |
| 146 | +impl<T: 'static> DerefMut for AllocatedThinBoxSlice<T> { |
| 147 | + #[inline] |
| 148 | + fn deref_mut(&mut self) -> &mut Self::Target { |
| 149 | + unsafe { slice::from_raw_parts_mut(self.ptr.as_ptr(), self.read_len()) } |
| 150 | + } |
| 151 | +} |
| 152 | + |
| 153 | +impl<T> AllocatedThinBoxSlice<MaybeUninit<T>> { |
| 154 | + #[inline] |
| 155 | + unsafe fn assume_init(self) -> AllocatedThinBoxSlice<T> { |
| 156 | + AllocatedThinBoxSlice { |
| 157 | + ptr: self.ptr.cast::<T>(), |
| 158 | + } |
| 159 | + } |
| 160 | +} |
| 161 | + |
| 162 | +impl<T: 'static> AllocatedThinBoxSlice<T> { |
| 163 | + #[inline] |
| 164 | + pub(super) fn run_drop(self) { |
| 165 | + unsafe { |
| 166 | + let len = self.read_len(); |
| 167 | + if len != 0 { |
| 168 | + let slice = ptr::slice_from_raw_parts_mut(self.ptr.as_ptr(), len); |
| 169 | + ptr::drop_in_place(slice); |
| 170 | + let alloc = self.ptr.cast::<usize>().as_ptr().sub(1); |
| 171 | + alloc::dealloc(alloc as *mut u8, Self::layout_for_len(len)); |
| 172 | + } |
| 173 | + } |
| 174 | + } |
| 175 | +} |
| 176 | + |
| 177 | +impl<T: 'static> Default for AllocatedThinBoxSlice<T> { |
| 178 | + #[inline] |
| 179 | + fn default() -> Self { |
| 180 | + AllocatedThinBoxSlice::empty() |
| 181 | + } |
| 182 | +} |
| 183 | + |
| 184 | +impl<T: Debug> Debug for AllocatedThinBoxSlice<T> { |
| 185 | + #[inline] |
| 186 | + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { |
| 187 | + <[T] as Debug>::fmt(&**self, f) |
| 188 | + } |
| 189 | +} |
| 190 | + |
| 191 | +impl<T: PartialEq> PartialEq for AllocatedThinBoxSlice<T> { |
| 192 | + #[inline] |
| 193 | + fn eq(&self, other: &Self) -> bool { |
| 194 | + <[T] as PartialEq>::eq(&**self, &**other) |
| 195 | + } |
| 196 | +} |
| 197 | + |
| 198 | +impl<T: Eq> Eq for AllocatedThinBoxSlice<T> {} |
| 199 | + |
| 200 | +impl<T: PartialOrd> PartialOrd for AllocatedThinBoxSlice<T> { |
| 201 | + #[inline] |
| 202 | + fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> { |
| 203 | + <[T] as PartialOrd>::partial_cmp(&**self, &**other) |
| 204 | + } |
| 205 | +} |
| 206 | + |
| 207 | +impl<T: Hash> Hash for AllocatedThinBoxSlice<T> { |
| 208 | + #[inline] |
| 209 | + fn hash<H: Hasher>(&self, state: &mut H) { |
| 210 | + <[T] as Hash>::hash(&**self, state) |
| 211 | + } |
| 212 | +} |
| 213 | + |
| 214 | +impl<T> FromIterator<T> for AllocatedThinBoxSlice<T> { |
| 215 | + fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self { |
| 216 | + let iter = iter.into_iter(); |
| 217 | + let (lower, upper) = iter.size_hint(); |
| 218 | + if Some(lower) == upper { |
| 219 | + let mut thin = AllocatedThinBoxSlice::<T>::new_uninit(lower); |
| 220 | + let mut i = 0; |
| 221 | + for item in iter { |
| 222 | + assert!(i < lower, "iterator produced more than promised"); |
| 223 | + MaybeUninit::write(&mut thin[i], item); |
| 224 | + i += 1; |
| 225 | + } |
| 226 | + assert_eq!(i, lower, "iterator produced less than promised"); |
| 227 | + unsafe { thin.assume_init() } |
| 228 | + } else { |
| 229 | + // TODO(nga): we can collect into partially initialized `ThinBoxSlice` |
| 230 | + // to get a chance of avoiding last reallocation. |
| 231 | + let vec = Vec::from_iter(iter); |
| 232 | + Self::from_iter(vec) |
| 233 | + } |
| 234 | + } |
| 235 | +} |
| 236 | + |
| 237 | +impl<T: Allocative> Allocative for AllocatedThinBoxSlice<T> { |
| 238 | + fn visit<'a, 'b: 'a>(&self, visitor: &'a mut allocative::Visitor<'b>) { |
| 239 | + let mut visitor = visitor.enter_self_sized::<Self>(); |
| 240 | + { |
| 241 | + let ptr_key = allocative::Key::new("ptr"); |
| 242 | + if self.len() == 0 { |
| 243 | + // Statically allocated data, so just report the pointer itself |
| 244 | + visitor.visit_simple(ptr_key, mem::size_of_val(&self.ptr)); |
| 245 | + } else { |
| 246 | + let mut visitor = |
| 247 | + visitor.enter_unique(allocative::Key::new("ptr"), mem::size_of_val(&self.ptr)); |
| 248 | + { |
| 249 | + let mut visitor = visitor.enter( |
| 250 | + allocative::Key::new("alloc"), |
| 251 | + Self::layout_for_len(self.len()).size(), |
| 252 | + ); |
| 253 | + visitor.visit_simple(allocative::Key::new("len"), mem::size_of::<usize>()); |
| 254 | + { |
| 255 | + let mut visitor = visitor |
| 256 | + .enter(allocative::Key::new("data"), mem::size_of_val::<[_]>(self)); |
| 257 | + visitor.visit_slice::<T>(self); |
| 258 | + visitor.exit(); |
| 259 | + } |
| 260 | + visitor.exit(); |
| 261 | + } |
| 262 | + visitor.exit(); |
| 263 | + } |
| 264 | + } |
| 265 | + visitor.exit(); |
| 266 | + } |
| 267 | +} |
| 268 | + |
| 269 | +#[cfg(test)] |
| 270 | +mod tests { |
| 271 | + use super::AllocatedThinBoxSlice; |
| 272 | + |
| 273 | + #[test] |
| 274 | + fn test_empty() { |
| 275 | + let thin = AllocatedThinBoxSlice::<String>::empty(); |
| 276 | + assert_eq!(0, thin.len()); |
| 277 | + thin.run_drop(); |
| 278 | + } |
| 279 | + |
| 280 | + #[test] |
| 281 | + fn test_from_iter_sized() { |
| 282 | + let thin = |
| 283 | + AllocatedThinBoxSlice::from_iter(["a".to_owned(), "bb".to_owned(), "ccc".to_owned()]); |
| 284 | + assert_eq!(["a".to_owned(), "bb".to_owned(), "ccc".to_owned()], *thin); |
| 285 | + thin.run_drop(); |
| 286 | + } |
| 287 | + |
| 288 | + #[test] |
| 289 | + fn test_from_iter_unknown_size() { |
| 290 | + let thin = AllocatedThinBoxSlice::from_iter( |
| 291 | + ["a".to_owned(), "b".to_owned(), "c".to_owned()] |
| 292 | + .into_iter() |
| 293 | + .filter(|_| true), |
| 294 | + ); |
| 295 | + assert_eq!(["a".to_owned(), "b".to_owned(), "c".to_owned()], *thin); |
| 296 | + thin.run_drop(); |
| 297 | + } |
| 298 | + |
| 299 | + /// If there are obvious memory violations, this test will catch them. |
| 300 | + #[test] |
| 301 | + fn test_stress() { |
| 302 | + for i in 0..1000 { |
| 303 | + let thin = AllocatedThinBoxSlice::from_iter((0..i).map(|j| j.to_string())); |
| 304 | + assert_eq!(i, thin.len()); |
| 305 | + assert_eq!((0..i).map(|j| j.to_string()).collect::<Vec<_>>(), *thin); |
| 306 | + thin.run_drop(); |
| 307 | + } |
| 308 | + } |
| 309 | +} |
0 commit comments