starlark/values/layout/
aligned_size.rs1use std::alloc::Layout;
19use std::mem;
20use std::ops::Add;
21use std::ops::Mul;
22use std::ops::Sub;
23use std::ptr::NonNull;
24
25use allocative::Allocative;
26use dupe::Dupe;
27
28use crate::values::layout::heap::repr::AValueHeader;
29
30#[derive(
32 Copy,
33 Clone,
34 Dupe,
35 Default,
36 Debug,
37 Eq,
38 PartialEq,
39 Ord,
40 PartialOrd,
41 Hash,
42 Allocative,
43 derive_more::Display
44)]
45#[repr(transparent)]
46pub(crate) struct AlignedSize {
47 bytes: u32,
49}
50
51impl AlignedSize {
52 pub(crate) const ZERO: AlignedSize = AlignedSize::new_bytes(0);
53
54 const MAX_SIZE: AlignedSize =
55 AlignedSize::new_bytes(u32::MAX as usize - AValueHeader::ALIGN + 1);
56
57 #[track_caller]
58 #[inline]
59 pub(crate) const fn new_bytes(bytes: usize) -> AlignedSize {
60 assert!(
61 bytes % AValueHeader::ALIGN == 0,
62 "AlignedSize must be aligned"
63 );
64 assert!(
65 bytes as u32 as usize == bytes,
66 "AlignedSize must not exceed u32::MAX"
67 );
68 let bytes = bytes as u32;
69 AlignedSize { bytes }
70 }
71
72 #[track_caller]
73 #[inline]
74 pub(crate) const fn align_up(bytes: usize) -> AlignedSize {
75 assert!(
76 bytes <= AlignedSize::MAX_SIZE.bytes() as usize,
77 "AlignedSize must not exceed u32::MAX"
78 );
79 let bytes = (bytes + AValueHeader::ALIGN - 1) & !(AValueHeader::ALIGN - 1);
80 let bytes = bytes as u32;
81 AlignedSize { bytes }
82 }
83
84 #[inline]
85 pub(crate) const fn bytes(self) -> u32 {
86 self.bytes
87 }
88
89 #[inline]
90 pub(crate) const fn of<T>() -> AlignedSize {
91 AlignedSize::align_up(mem::size_of::<T>())
92 }
93
94 #[inline]
95 pub(crate) const fn layout(self) -> Layout {
96 match Layout::from_size_align(self.bytes as usize, AValueHeader::ALIGN) {
97 Ok(layout) => layout,
98 Err(_) => panic!("Layout::from_size_align failed"),
99 }
100 }
101
102 #[inline]
103 pub(crate) fn checked_next_power_of_two(self) -> Option<AlignedSize> {
104 let bytes = self.bytes.checked_next_power_of_two()?;
105 Some(AlignedSize::new_bytes(bytes as usize))
106 }
107
108 #[inline]
109 pub(crate) fn unchecked_sub(self, rhs: AlignedSize) -> AlignedSize {
110 debug_assert!(self.bytes >= rhs.bytes, "{:?} - {:?}", self, rhs);
111 AlignedSize {
112 bytes: self.bytes - rhs.bytes,
113 }
114 }
115
116 #[inline]
117 pub(crate) fn ptr_diff(begin: NonNull<usize>, end: NonNull<usize>) -> AlignedSize {
118 unsafe { AlignedSize::new_bytes(end.as_ptr().byte_offset_from(begin.as_ptr()) as usize) }
119 }
120}
121
122impl Add for AlignedSize {
123 type Output = AlignedSize;
124
125 #[track_caller]
126 #[inline]
127 fn add(self, rhs: AlignedSize) -> AlignedSize {
128 let bytes = self.bytes.checked_add(rhs.bytes).unwrap();
129 AlignedSize { bytes }
130 }
131}
132
133impl Sub for AlignedSize {
134 type Output = AlignedSize;
135
136 #[track_caller]
137 #[inline]
138 fn sub(self, rhs: AlignedSize) -> AlignedSize {
139 let bytes = self.bytes.checked_sub(rhs.bytes).unwrap();
140 AlignedSize { bytes }
141 }
142}
143
144impl Mul<u32> for AlignedSize {
145 type Output = AlignedSize;
146
147 #[track_caller]
148 #[inline]
149 fn mul(self, rhs: u32) -> Self::Output {
150 let bytes = self.bytes.checked_mul(rhs).unwrap();
151 AlignedSize { bytes }
152 }
153}
154
155#[cfg(test)]
156mod tests {
157 use crate::values::layout::aligned_size::AlignedSize;
158 use crate::values::layout::heap::repr::AValueHeader;
159
160 #[test]
161 fn test_checked_next_power_of_two() {
162 assert_eq!(
163 AlignedSize::new_bytes(AValueHeader::ALIGN),
164 AlignedSize::new_bytes(AValueHeader::ALIGN)
165 .checked_next_power_of_two()
166 .unwrap()
167 );
168 assert_eq!(
169 AlignedSize::new_bytes(2 * AValueHeader::ALIGN),
170 AlignedSize::new_bytes(2 * AValueHeader::ALIGN)
171 .checked_next_power_of_two()
172 .unwrap()
173 );
174 assert_eq!(
175 AlignedSize::new_bytes(4 * AValueHeader::ALIGN),
176 AlignedSize::new_bytes(3 * AValueHeader::ALIGN)
177 .checked_next_power_of_two()
178 .unwrap()
179 );
180 assert_eq!(
181 AlignedSize::new_bytes(8 * AValueHeader::ALIGN),
182 AlignedSize::new_bytes(5 * AValueHeader::ALIGN)
183 .checked_next_power_of_two()
184 .unwrap()
185 );
186 }
187
188 #[test]
189 fn test_sub() {
190 assert_eq!(
191 AlignedSize::new_bytes(2 * AValueHeader::ALIGN),
192 AlignedSize::new_bytes(5 * AValueHeader::ALIGN)
193 - AlignedSize::new_bytes(3 * AValueHeader::ALIGN)
194 );
195 }
196}