Line data Source code
1 : use core::slice;
2 : use std::{
3 : alloc::{self, Layout},
4 : cmp,
5 : mem::ManuallyDrop,
6 : };
7 :
8 : use super::alignment::{Alignment, ConstAlign};
9 :
10 : #[derive(Debug)]
11 : struct AlignedBufferPtr(*mut u8);
12 :
13 : // SAFETY: We gurantees no one besides `IoBufferPtr` itself has the raw pointer.
14 : unsafe impl Send for AlignedBufferPtr {}
15 :
16 : // SAFETY: We gurantees no one besides `IoBufferPtr` itself has the raw pointer.
17 : unsafe impl Sync for AlignedBufferPtr {}
18 :
19 : /// An aligned buffer type.
20 : #[derive(Debug)]
21 : pub struct RawAlignedBuffer<A: Alignment> {
22 : ptr: AlignedBufferPtr,
23 : capacity: usize,
24 : len: usize,
25 : align: A,
26 : }
27 :
28 : impl<const A: usize> RawAlignedBuffer<ConstAlign<A>> {
29 : /// Constructs a new, empty `IoBufferMut` with at least the specified capacity and alignment.
30 : ///
31 : /// The buffer will be able to hold at most `capacity` elements and will never resize.
32 : ///
33 : ///
34 : /// # Panics
35 : ///
36 : /// Panics if the new capacity exceeds `isize::MAX` _bytes_, or if the following alignment requirement is not met:
37 : /// * `align` must not be zero,
38 : ///
39 : /// * `align` must be a power of two,
40 : ///
41 : /// * `capacity`, when rounded up to the nearest multiple of `align`,
42 : /// must not overflow isize (i.e., the rounded value must be
43 : /// less than or equal to `isize::MAX`).
44 996205 : pub fn with_capacity(capacity: usize) -> Self {
45 996205 : let align = ConstAlign::<A>;
46 996205 : let layout = Layout::from_size_align(capacity, align.align()).expect("Invalid layout");
47 :
48 : // SAFETY: Making an allocation with a sized and aligned layout. The memory is manually freed with the same layout.
49 996205 : let ptr = unsafe {
50 996205 : let ptr = alloc::alloc(layout);
51 996205 : if ptr.is_null() {
52 0 : alloc::handle_alloc_error(layout);
53 996205 : }
54 996205 : AlignedBufferPtr(ptr)
55 996205 : };
56 996205 :
57 996205 : RawAlignedBuffer {
58 996205 : ptr,
59 996205 : capacity,
60 996205 : len: 0,
61 996205 : align,
62 996205 : }
63 996205 : }
64 : }
65 :
66 : impl<A: Alignment> RawAlignedBuffer<A> {
67 : /// Returns the total number of bytes the buffer can hold.
68 : #[inline]
69 155383125 : pub fn capacity(&self) -> usize {
70 155383125 : self.capacity
71 155383125 : }
72 :
73 : /// Returns the alignment of the buffer.
74 : #[inline]
75 30 : pub fn align(&self) -> usize {
76 30 : self.align.align()
77 30 : }
78 :
79 : /// Returns the number of bytes in the buffer, also referred to as its 'length'.
80 : #[inline]
81 161575840 : pub fn len(&self) -> usize {
82 161575840 : self.len
83 161575840 : }
84 :
85 : /// Force the length of the buffer to `new_len`.
86 : #[inline]
87 38593757 : pub unsafe fn set_len(&mut self, new_len: usize) {
88 38593757 : debug_assert!(new_len <= self.capacity());
89 38593757 : self.len = new_len;
90 38593757 : }
91 :
92 : #[inline]
93 38770824 : pub fn as_ptr(&self) -> *const u8 {
94 38770824 : self.ptr.0
95 38770824 : }
96 :
97 : #[inline]
98 40578433 : pub fn as_mut_ptr(&mut self) -> *mut u8 {
99 40578433 : self.ptr.0
100 40578433 : }
101 :
102 : /// Extracts a slice containing the entire buffer.
103 : ///
104 : /// Equivalent to `&s[..]`.
105 : #[inline]
106 3834287 : pub fn as_slice(&self) -> &[u8] {
107 3834287 : // SAFETY: The pointer is valid and `len` bytes are initialized.
108 3834287 : unsafe { slice::from_raw_parts(self.as_ptr(), self.len) }
109 3834287 : }
110 :
111 : /// Extracts a mutable slice of the entire buffer.
112 : ///
113 : /// Equivalent to `&mut s[..]`.
114 0 : pub fn as_mut_slice(&mut self) -> &mut [u8] {
115 0 : // SAFETY: The pointer is valid and `len` bytes are initialized.
116 0 : unsafe { slice::from_raw_parts_mut(self.as_mut_ptr(), self.len) }
117 0 : }
118 :
119 : /// Drops the all the contents of the buffer, setting its length to `0`.
120 : #[inline]
121 26 : pub fn clear(&mut self) {
122 26 : self.len = 0;
123 26 : }
124 :
125 : /// Reserves capacity for at least `additional` more bytes to be inserted
126 : /// in the given `IoBufferMut`. The collection may reserve more space to
127 : /// speculatively avoid frequent reallocations. After calling `reserve`,
128 : /// capacity will be greater than or equal to `self.len() + additional`.
129 : /// Does nothing if capacity is already sufficient.
130 : ///
131 : /// # Panics
132 : ///
133 : /// Panics if the new capacity exceeds `isize::MAX` _bytes_.
134 22 : pub fn reserve(&mut self, additional: usize) {
135 22 : if additional > self.capacity() - self.len() {
136 6 : self.reserve_inner(additional);
137 16 : }
138 22 : }
139 :
140 6 : fn reserve_inner(&mut self, additional: usize) {
141 6 : let Some(required_cap) = self.len().checked_add(additional) else {
142 0 : capacity_overflow()
143 : };
144 :
145 6 : let old_capacity = self.capacity();
146 6 : let align = self.align();
147 6 : // This guarantees exponential growth. The doubling cannot overflow
148 6 : // because `cap <= isize::MAX` and the type of `cap` is `usize`.
149 6 : let cap = cmp::max(old_capacity * 2, required_cap);
150 6 :
151 6 : if !is_valid_alloc(cap) {
152 0 : capacity_overflow()
153 6 : }
154 6 : let new_layout = Layout::from_size_align(cap, self.align()).expect("Invalid layout");
155 6 :
156 6 : let old_ptr = self.as_mut_ptr();
157 :
158 : // SAFETY: old allocation was allocated with std::alloc::alloc with the same layout,
159 : // and we panics on null pointer.
160 6 : let (ptr, cap) = unsafe {
161 6 : let old_layout = Layout::from_size_align_unchecked(old_capacity, align);
162 6 : let ptr = alloc::realloc(old_ptr, old_layout, new_layout.size());
163 6 : if ptr.is_null() {
164 0 : alloc::handle_alloc_error(new_layout);
165 6 : }
166 6 : (AlignedBufferPtr(ptr), cap)
167 6 : };
168 6 :
169 6 : self.ptr = ptr;
170 6 : self.capacity = cap;
171 6 : }
172 :
173 : /// Shortens the buffer, keeping the first len bytes.
174 10 : pub fn truncate(&mut self, len: usize) {
175 10 : if len > self.len {
176 0 : return;
177 10 : }
178 10 : self.len = len;
179 10 : }
180 :
181 : /// Consumes and leaks the `IoBufferMut`, returning a mutable reference to the contents, &'a mut [u8].
182 88 : pub fn leak<'a>(self) -> &'a mut [u8] {
183 88 : let mut buf = ManuallyDrop::new(self);
184 88 : // SAFETY: leaking the buffer as intended.
185 88 : unsafe { slice::from_raw_parts_mut(buf.as_mut_ptr(), buf.len) }
186 88 : }
187 : }
188 :
189 0 : fn capacity_overflow() -> ! {
190 0 : panic!("capacity overflow")
191 : }
192 :
193 : // We need to guarantee the following:
194 : // * We don't ever allocate `> isize::MAX` byte-size objects.
195 : // * We don't overflow `usize::MAX` and actually allocate too little.
196 : //
197 : // On 64-bit we just need to check for overflow since trying to allocate
198 : // `> isize::MAX` bytes will surely fail. On 32-bit and 16-bit we need to add
199 : // an extra guard for this in case we're running on a platform which can use
200 : // all 4GB in user-space, e.g., PAE or x32.
201 : #[inline]
202 6 : fn is_valid_alloc(alloc_size: usize) -> bool {
203 6 : !(usize::BITS < 64 && alloc_size > isize::MAX as usize)
204 6 : }
205 :
206 : impl<A: Alignment> Drop for RawAlignedBuffer<A> {
207 996117 : fn drop(&mut self) {
208 996117 : // SAFETY: memory was allocated with std::alloc::alloc with the same layout.
209 996117 : unsafe {
210 996117 : alloc::dealloc(
211 996117 : self.as_mut_ptr(),
212 996117 : Layout::from_size_align_unchecked(self.capacity, self.align.align()),
213 996117 : )
214 996117 : }
215 996117 : }
216 : }
|
