1// Copyright (c) 2017-2022, The rav1e contributors. All rights reserved
2//
3// This source code is subject to the terms of the BSD 2 Clause License and
4// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
5// was not distributed with this source code in the LICENSE file, you can
6// obtain it at www.aomedia.org/license/software. If the Alliance for Open
7// Media Patent License 1.0 was not distributed with this source code in the
8// PATENTS file, you can obtain it at www.aomedia.org/license/patent.
9
10use std::alloc::{alloc, dealloc, Layout};
11use std::mem::MaybeUninit;
12use std::ptr;
13use std::{fmt, mem};
14
15#[repr(align(64))]
16pub struct Align64;
17
18// A 64 byte aligned piece of data.
19// # Examples
20// ```
21// let mut x: Aligned<[i16; 64 * 64]> = Aligned::new([0; 64 * 64]);
22// assert!(x.data.as_ptr() as usize % 16 == 0);
23//
24// let mut x: Aligned<[i16; 64 * 64]> = Aligned::uninitialized();
25// assert!(x.data.as_ptr() as usize % 16 == 0);
26// ```
27pub struct Aligned<T> {
28 _alignment: [Align64; 0],
29 pub data: T,
30}
31
32#[cfg(any(test, feature = "bench"))]
33impl<const N: usize, T> Aligned<[T; N]> {
34 #[inline(always)]
35 pub fn from_fn<F>(cb: F) -> Self
36 where
37 F: FnMut(usize) -> T,
38 {
39 Aligned { _alignment: [], data: std::array::from_fn(cb) }
40 }
41}
42
43impl<const N: usize, T> Aligned<[MaybeUninit<T>; N]> {
44 #[inline(always)]
45 pub const fn uninit_array() -> Self {
46 Aligned {
47 _alignment: [],
48 // SAFETY: Uninitialized [MaybeUninit<T>; N] is valid.
49 data: unsafe { MaybeUninit::uninit().assume_init() },
50 }
51 }
52}
53
54impl<T> Aligned<T> {
55 pub const fn new(data: T) -> Self {
56 Aligned { _alignment: [], data }
57 }
58 #[allow(clippy::uninit_assumed_init)]
59 /// # Safety
60 ///
61 /// The resulting `Aligned<T>` *must* be written to before it is read from.
62 pub const unsafe fn uninitialized() -> Self {
63 Self::new(data:MaybeUninit::uninit().assume_init())
64 }
65}
66
67/// An analog to a Box<[T]> where the underlying slice is aligned.
68/// Alignment is according to the architecture-specific SIMD constraints.
69pub struct AlignedBoxedSlice<T> {
70 ptr: std::ptr::NonNull<T>,
71 len: usize,
72}
73
74impl<T> AlignedBoxedSlice<T> {
75 // Data alignment in bytes.
76 cfg_if::cfg_if! {
77 if #[cfg(target_arch = "wasm32")] {
78 // FIXME: wasm32 allocator fails for alignment larger than 3
79 const DATA_ALIGNMENT_LOG2: usize = 3;
80 } else {
81 const DATA_ALIGNMENT_LOG2: usize = 6;
82 }
83 }
84
85 const fn layout(len: usize) -> Layout {
86 // SAFETY: We are ensuring that `align` is non-zero and is a multiple of 2.
87 unsafe {
88 Layout::from_size_align_unchecked(
89 len * mem::size_of::<T>(),
90 1 << Self::DATA_ALIGNMENT_LOG2,
91 )
92 }
93 }
94
95 fn alloc(len: usize) -> std::ptr::NonNull<T> {
96 // SAFETY: We are not calling this with a null pointer, so it's safe.
97 unsafe { ptr::NonNull::new_unchecked(alloc(Self::layout(len)) as *mut T) }
98 }
99
100 /// Creates a [`AlignedBoxedSlice`] with a slice of length [`len`] filled with
101 /// [`val`].
102 pub fn new(len: usize, val: T) -> Self
103 where
104 T: Clone,
105 {
106 let mut output = Self { ptr: Self::alloc(len), len };
107
108 for a in output.iter_mut() {
109 *a = val.clone();
110 }
111
112 output
113 }
114}
115
116impl<T: fmt::Debug> fmt::Debug for AlignedBoxedSlice<T> {
117 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
118 fmt::Debug::fmt(&**self, f)
119 }
120}
121
122impl<T> std::ops::Deref for AlignedBoxedSlice<T> {
123 type Target = [T];
124
125 fn deref(&self) -> &[T] {
126 // SAFETY: We know that `self.ptr` is not null, and we know its length.
127 unsafe {
128 let p: *mut T = self.ptr.as_ptr();
129
130 std::slice::from_raw_parts(data:p, self.len)
131 }
132 }
133}
134
135impl<T> std::ops::DerefMut for AlignedBoxedSlice<T> {
136 fn deref_mut(&mut self) -> &mut [T] {
137 // SAFETY: We know that `self.ptr` is not null, and we know its length.
138 unsafe {
139 let p: *mut T = self.ptr.as_ptr();
140
141 std::slice::from_raw_parts_mut(data:p, self.len)
142 }
143 }
144}
145
146impl<T> std::ops::Drop for AlignedBoxedSlice<T> {
147 fn drop(&mut self) {
148 // SAFETY: We know that the contents of this struct are aligned and valid to drop.
149 unsafe {
150 for a: &mut T in self.iter_mut() {
151 ptr::drop_in_place(to_drop:a)
152 }
153
154 dealloc(self.ptr.as_ptr() as *mut u8, Self::layout(self.len));
155 }
156 }
157}
158
159unsafe impl<T> Send for AlignedBoxedSlice<T> where T: Send {}
160unsafe impl<T> Sync for AlignedBoxedSlice<T> where T: Sync {}
161
162#[cfg(test)]
163mod test {
164 use super::*;
165
166 fn is_aligned<T>(ptr: *const T, n: usize) -> bool {
167 ((ptr as usize) & ((1 << n) - 1)) == 0
168 }
169
170 #[test]
171 fn sanity_stack() {
172 let a: Aligned<_> = Aligned::new([0u8; 3]);
173 assert!(is_aligned(a.data.as_ptr(), 4));
174 }
175
176 #[test]
177 fn sanity_heap() {
178 let a: AlignedBoxedSlice<_> = AlignedBoxedSlice::new(3, 0u8);
179 assert!(is_aligned(a.as_ptr(), 4));
180 }
181}
182