starlark/values/layout/
aligned_size.rs

1/*
2 * Copyright 2019 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
18use 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/// Allocations in Starlark are word-aligned, and this type represents the size of an allocation.
31#[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    /// Starlark only supports objects smaller than 1<<32.
48    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}