rav1e/util/
align.rs

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(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 = self.ptr.as_ptr();
129
130      std::slice::from_raw_parts(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 = self.ptr.as_ptr();
140
141      std::slice::from_raw_parts_mut(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 in self.iter_mut() {
151        ptr::drop_in_place(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}
rav1e/util/align.rs

rav1e/util/
align.rs
