Skip to content

Commit 27fcdb8

Browse files
mbrubeckmbrubeck
committed
Specialize Vec::from_elem<u8> to use calloc or memset
Fixes rust-lang#38723.
1 parent a559452 commit 27fcdb8

File tree

6 files changed

+137
-8
lines changed

6 files changed

+137
-8
lines changed

‎src/doc/unstable-book/src/allocator.md

+5
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,11 @@ pub extern fn __rust_allocate(size: usize, _align: usize) -> *mut u8 {
5151
unsafe { libc::malloc(size as libc::size_t) as *mut u8 }
5252
}
5353
54+
#[no_mangle]
55+
pub extern fn __rust_allocate_zeroed(size: usize, _align: usize) -> *mut u8 {
56+
unsafe { libc::calloc(size as libc::size_t, 1) as *mut u8 }
57+
}
58+
5459
#[no_mangle]
5560
pub extern fn __rust_deallocate(ptr: *mut u8, _old_size: usize, _align: usize) {
5661
unsafe { libc::free(ptr as *mut libc::c_void) }

‎src/liballoc/heap.rs

+32
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ use core::intrinsics::{min_align_of_val, size_of_val};
2323
extern "C" {
2424
#[allocator]
2525
fn __rust_allocate(size: usize, align: usize) -> *mut u8;
26+
fn __rust_allocate_zeroed(size: usize, align: usize) -> *mut u8;
2627
fn __rust_deallocate(ptr: *mut u8, old_size: usize, align: usize);
2728
fn __rust_reallocate(ptr: *mut u8, old_size: usize, size: usize, align: usize) -> *mut u8;
2829
fn __rust_reallocate_inplace(ptr: *mut u8,
@@ -59,6 +60,20 @@ pub unsafe fn allocate(size: usize, align: usize) -> *mut u8 {
5960
__rust_allocate(size, align)
6061
}
6162

63+
/// Return a pointer to `size` bytes of memory aligned to `align` and
64+
/// initialized to zeroes.
65+
///
66+
/// On failure, return a null pointer.
67+
///
68+
/// Behavior is undefined if the requested size is 0 or the alignment is not a
69+
/// power of 2. The alignment must be no larger than the largest supported page
70+
/// size on the platform.
71+
#[inline]
72+
pub unsafe fn allocate_zeroed(size: usize, align: usize) -> *mut u8 {
73+
check_size_and_alignment(size, align);
74+
__rust_allocate_zeroed(size, align)
75+
}
76+
6277
/// Resize the allocation referenced by `ptr` to `size` bytes.
6378
///
6479
/// On failure, return a null pointer and leave the original allocation intact.
@@ -162,6 +177,23 @@ mod tests {
162177
use boxed::Box;
163178
use heap;
164179

180+
#[test]
181+
fn allocate_zeroed() {
182+
unsafe {
183+
let size = 1024;
184+
let mut ptr = heap::allocate_zeroed(size, 1);
185+
if ptr.is_null() {
186+
::oom()
187+
}
188+
let end = ptr.offset(size as isize);
189+
while ptr < end {
190+
assert_eq!(*ptr, 0);
191+
ptr = ptr.offset(1);
192+
}
193+
heap::deallocate(ptr, size, 1);
194+
}
195+
}
196+
165197
#[test]
166198
fn basic_reallocate_inplace_noop() {
167199
unsafe {

‎src/liballoc/raw_vec.rs

+15-1
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,16 @@ impl<T> RawVec<T> {
8282
///
8383
/// Aborts on OOM
8484
pub fn with_capacity(cap: usize) -> Self {
85+
RawVec::allocate(cap, false)
86+
}
87+
88+
/// Like `with_capacity` but guarantees the buffer is zeroed.
89+
pub fn with_capacity_zeroed(cap: usize) -> Self {
90+
RawVec::allocate(cap, true)
91+
}
92+
93+
#[inline]
94+
fn allocate(cap: usize, zeroed: bool) -> Self {
8595
unsafe {
8696
let elem_size = mem::size_of::<T>();
8797

@@ -93,7 +103,11 @@ impl<T> RawVec<T> {
93103
heap::EMPTY as *mut u8
94104
} else {
95105
let align = mem::align_of::<T>();
96-
let ptr = heap::allocate(alloc_size, align);
106+
let ptr = if zeroed {
107+
heap::allocate_zeroed(alloc_size, align)
108+
} else {
109+
heap::allocate(alloc_size, align)
110+
};
97111
if ptr.is_null() {
98112
oom()
99113
}

‎src/liballoc_jemalloc/lib.rs

+21
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,10 @@ mod imp {
3838
target_os = "dragonfly", target_os = "windows"),
3939
link_name = "je_mallocx")]
4040
fn mallocx(size: size_t, flags: c_int) -> *mut c_void;
41+
#[cfg_attr(any(target_os = "macos", target_os = "android", target_os = "ios",
42+
target_os = "dragonfly", target_os = "windows"),
43+
link_name = "je_calloc")]
44+
fn calloc(size: size_t, flags: c_int) -> *mut c_void;
4145
#[cfg_attr(any(target_os = "macos", target_os = "android", target_os = "ios",
4246
target_os = "dragonfly", target_os = "windows"),
4347
link_name = "je_rallocx")]
@@ -56,6 +60,8 @@ mod imp {
5660
fn nallocx(size: size_t, flags: c_int) -> size_t;
5761
}
5862

63+
const MALLOCX_ZERO: c_int = 0x40;
64+
5965
// The minimum alignment guaranteed by the architecture. This value is used to
6066
// add fast paths for low alignment values. In practice, the alignment is a
6167
// constant at the call site and the branch will be optimized out.
@@ -91,6 +97,16 @@ mod imp {
9197
unsafe { mallocx(size as size_t, flags) as *mut u8 }
9298
}
9399

100+
#[no_mangle]
101+
pub extern "C" fn __rust_allocate_zeroed(size: usize, align: usize) -> *mut u8 {
102+
if align <= MIN_ALIGN {
103+
unsafe { calloc(size as size_t, 1) as *mut u8 }
104+
} else {
105+
let flags = align_to_flags(align) | MALLOCX_ZERO;
106+
unsafe { mallocx(size as size_t, flags) as *mut u8 }
107+
}
108+
}
109+
94110
#[no_mangle]
95111
pub extern "C" fn __rust_reallocate(ptr: *mut u8,
96112
_old_size: usize,
@@ -135,6 +151,11 @@ mod imp {
135151
bogus()
136152
}
137153

154+
#[no_mangle]
155+
pub extern "C" fn __rust_allocate_zeroed(_size: usize, _align: usize) -> *mut u8 {
156+
bogus()
157+
}
158+
138159
#[no_mangle]
139160
pub extern "C" fn __rust_reallocate(_ptr: *mut u8,
140161
_old_size: usize,

‎src/liballoc_system/lib.rs

+32-4
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
1+
// Copyright 2016 The Rust Project Developers. See the COPYRIGHT
22
// file at the top-level directory of this distribution and at
33
// http://rust-lang.org/COPYRIGHT.
44
//
@@ -44,6 +44,11 @@ pub extern "C" fn __rust_allocate(size: usize, align: usize) -> *mut u8 {
4444
unsafe { imp::allocate(size, align) }
4545
}
4646

47+
#[no_mangle]
48+
pub extern "C" fn __rust_allocate_zeroed(size: usize, align: usize) -> *mut u8 {
49+
unsafe { imp::allocate_zeroed(size, align) }
50+
}
51+
4752
#[no_mangle]
4853
pub extern "C" fn __rust_deallocate(ptr: *mut u8, old_size: usize, align: usize) {
4954
unsafe { imp::deallocate(ptr, old_size, align) }
@@ -121,6 +126,18 @@ mod imp {
121126
}
122127
}
123128

129+
pub unsafe fn allocate_zeroed(size: usize, align: usize) -> *mut u8 {
130+
if align <= MIN_ALIGN {
131+
libc::calloc(size as libc::size_t, 1) as *mut u8
132+
} else {
133+
let ptr = aligned_malloc(size, align);
134+
if !ptr.is_null() {
135+
ptr::write_bytes(ptr, 0, size);
136+
}
137+
ptr
138+
}
139+
}
140+
124141
pub unsafe fn reallocate(ptr: *mut u8, old_size: usize, size: usize, align: usize) -> *mut u8 {
125142
if align <= MIN_ALIGN {
126143
libc::realloc(ptr as *mut libc::c_void, size as libc::size_t) as *mut u8
@@ -173,6 +190,8 @@ mod imp {
173190
#[repr(C)]
174191
struct Header(*mut u8);
175192

193+
194+
const HEAP_ZERO_MEMORY: DWORD = 0x00000008;
176195
const HEAP_REALLOC_IN_PLACE_ONLY: DWORD = 0x00000010;
177196

178197
unsafe fn get_header<'a>(ptr: *mut u8) -> &'a mut Header {
@@ -185,18 +204,27 @@ mod imp {
185204
aligned
186205
}
187206

188-
pub unsafe fn allocate(size: usize, align: usize) -> *mut u8 {
207+
#[inline]
208+
unsafe fn allocate_with_flags(size: usize, align: usize, flags: DWORD) -> *mut u8 {
189209
if align <= MIN_ALIGN {
190-
HeapAlloc(GetProcessHeap(), 0, size as SIZE_T) as *mut u8
210+
HeapAlloc(GetProcessHeap(), flags, size as SIZE_T) as *mut u8
191211
} else {
192-
let ptr = HeapAlloc(GetProcessHeap(), 0, (size + align) as SIZE_T) as *mut u8;
212+
let ptr = HeapAlloc(GetProcessHeap(), flags, (size + align) as SIZE_T) as *mut u8;
193213
if ptr.is_null() {
194214
return ptr;
195215
}
196216
align_ptr(ptr, align)
197217
}
198218
}
199219

220+
pub unsafe fn allocate(size: usize, align: usize) -> *mut u8 {
221+
allocate_with_flags(size, align, 0)
222+
}
223+
224+
pub unsafe fn allocate_zeroed(size: usize, align: usize) -> *mut u8 {
225+
allocate_with_flags(size, align, HEAP_ZERO_MEMORY)
226+
}
227+
200228
pub unsafe fn reallocate(ptr: *mut u8, _old_size: usize, size: usize, align: usize) -> *mut u8 {
201229
if align <= MIN_ALIGN {
202230
HeapReAlloc(GetProcessHeap(), 0, ptr as LPVOID, size as SIZE_T) as *mut u8

‎src/libcollections/vec.rs

+32-3
Original file line numberDiff line numberDiff line change
@@ -1369,9 +1369,38 @@ impl<T: PartialEq> Vec<T> {
13691369
#[doc(hidden)]
13701370
#[stable(feature = "rust1", since = "1.0.0")]
13711371
pub fn from_elem<T: Clone>(elem: T, n: usize) -> Vec<T> {
1372-
let mut v = Vec::with_capacity(n);
1373-
v.extend_with_element(n, elem);
1374-
v
1372+
<T as SpecFromElem>::from_elem(elem, n)
1373+
}
1374+
1375+
// Specialization trait used for Vec::from_elem
1376+
trait SpecFromElem: Sized {
1377+
fn from_elem(elem: Self, n: usize) -> Vec<Self>;
1378+
}
1379+
1380+
impl<T: Clone> SpecFromElem for T {
1381+
default fn from_elem(elem: Self, n: usize) -> Vec<Self> {
1382+
let mut v = Vec::with_capacity(n);
1383+
v.extend_with_element(n, elem);
1384+
v
1385+
}
1386+
}
1387+
1388+
impl SpecFromElem for u8 {
1389+
#[inline]
1390+
fn from_elem(elem: u8, n: usize) -> Vec<u8> {
1391+
if elem == 0 {
1392+
return Vec {
1393+
buf: RawVec::with_capacity_zeroed(n),
1394+
len: n,
1395+
}
1396+
}
1397+
unsafe {
1398+
let mut v = Vec::with_capacity(n);
1399+
ptr::write_bytes(v.as_mut_ptr(), elem, n);
1400+
v.set_len(n);
1401+
v
1402+
}
1403+
}
13751404
}
13761405

13771406
////////////////////////////////////////////////////////////////////////////////

0 commit comments

Comments
 (0)