Skip to content

Commit 7627e3d

Browse files
borsbors
committed
Auto merge of #40409 - mbrubeck:calloc, r=sfackler
Specialize Vec::from_elem to use calloc Fixes #38723. This specializes the implementation for `u8` only, but it could be extended to other zeroable types if desired. I haven't tested this extensively, but I did verify that it gives the expected performance boost for large `vec![0; n]` allocations with both alloc_system and jemalloc, on Linux. (I have not tested or even built the Windows code.)
2 parents 9af79f5 + aad2062 commit 7627e3d

File tree

8 files changed

+183
-7
lines changed

8 files changed

+183
-7
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

+34
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,25 @@ mod tests {
162177
use boxed::Box;
163178
use heap;
164179

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

‎src/liballoc/raw_vec.rs

+16-1
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,18 @@ impl<T> RawVec<T> {
8181
/// # Aborts
8282
///
8383
/// Aborts on OOM
84+
#[inline]
8485
pub fn with_capacity(cap: usize) -> Self {
86+
RawVec::allocate(cap, false)
87+
}
88+
89+
/// Like `with_capacity` but guarantees the buffer is zeroed.
90+
#[inline]
91+
pub fn with_capacity_zeroed(cap: usize) -> Self {
92+
RawVec::allocate(cap, true)
93+
}
94+
95+
fn allocate(cap: usize, zeroed: bool) -> Self {
8596
unsafe {
8697
let elem_size = mem::size_of::<T>();
8798

@@ -93,7 +104,11 @@ impl<T> RawVec<T> {
93104
heap::EMPTY as *mut u8
94105
} else {
95106
let align = mem::align_of::<T>();
96-
let ptr = heap::allocate(alloc_size, align);
107+
let ptr = if zeroed {
108+
heap::allocate_zeroed(alloc_size, align)
109+
} else {
110+
heap::allocate(alloc_size, align)
111+
};
97112
if ptr.is_null() {
98113
oom()
99114
}

‎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", target_env = "musl"),
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", target_env = "musl"),
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", target_env = "musl"),
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

+31-3
Original file line numberDiff line numberDiff line change
@@ -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/lib.rs

+2
Original file line numberDiff line numberDiff line change
@@ -35,13 +35,15 @@
3535
#![feature(box_patterns)]
3636
#![feature(box_syntax)]
3737
#![cfg_attr(not(test), feature(char_escape_debug))]
38+
#![cfg_attr(not(test), feature(core_float))]
3839
#![feature(core_intrinsics)]
3940
#![feature(dropck_eyepatch)]
4041
#![feature(exact_size_is_empty)]
4142
#![feature(fmt_internals)]
4243
#![feature(fused)]
4344
#![feature(generic_param_attrs)]
4445
#![feature(heap_api)]
46+
#![feature(i128_type)]
4547
#![feature(inclusive_range)]
4648
#![feature(lang_items)]
4749
#![feature(manually_drop)]

‎src/libcollections/vec.rs

+69-3
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,8 @@ use core::hash::{self, Hash};
7777
use core::intrinsics::{arith_offset, assume};
7878
use core::iter::{FromIterator, FusedIterator, TrustedLen};
7979
use core::mem;
80+
#[cfg(not(test))]
81+
use core::num::Float;
8082
use core::ops::{InPlace, Index, IndexMut, Place, Placer};
8183
use core::ops;
8284
use core::ptr;
@@ -1370,11 +1372,75 @@ impl<T: PartialEq> Vec<T> {
13701372
#[doc(hidden)]
13711373
#[stable(feature = "rust1", since = "1.0.0")]
13721374
pub fn from_elem<T: Clone>(elem: T, n: usize) -> Vec<T> {
1373-
let mut v = Vec::with_capacity(n);
1374-
v.extend_with_element(n, elem);
1375-
v
1375+
<T as SpecFromElem>::from_elem(elem, n)
1376+
}
1377+
1378+
// Specialization trait used for Vec::from_elem
1379+
trait SpecFromElem: Sized {
1380+
fn from_elem(elem: Self, n: usize) -> Vec<Self>;
13761381
}
13771382

1383+
impl<T: Clone> SpecFromElem for T {
1384+
default fn from_elem(elem: Self, n: usize) -> Vec<Self> {
1385+
let mut v = Vec::with_capacity(n);
1386+
v.extend_with_element(n, elem);
1387+
v
1388+
}
1389+
}
1390+
1391+
impl SpecFromElem for u8 {
1392+
#[inline]
1393+
fn from_elem(elem: u8, n: usize) -> Vec<u8> {
1394+
if elem == 0 {
1395+
return Vec {
1396+
buf: RawVec::with_capacity_zeroed(n),
1397+
len: n,
1398+
}
1399+
}
1400+
unsafe {
1401+
let mut v = Vec::with_capacity(n);
1402+
ptr::write_bytes(v.as_mut_ptr(), elem, n);
1403+
v.set_len(n);
1404+
v
1405+
}
1406+
}
1407+
}
1408+
1409+
macro_rules! impl_spec_from_elem {
1410+
($t: ty, $is_zero: expr) => {
1411+
impl SpecFromElem for $t {
1412+
#[inline]
1413+
fn from_elem(elem: $t, n: usize) -> Vec<$t> {
1414+
if $is_zero(elem) {
1415+
return Vec {
1416+
buf: RawVec::with_capacity_zeroed(n),
1417+
len: n,
1418+
}
1419+
}
1420+
let mut v = Vec::with_capacity(n);
1421+
v.extend_with_element(n, elem);
1422+
v
1423+
}
1424+
}
1425+
};
1426+
}
1427+
1428+
impl_spec_from_elem!(i8, |x| x == 0);
1429+
impl_spec_from_elem!(i16, |x| x == 0);
1430+
impl_spec_from_elem!(i32, |x| x == 0);
1431+
impl_spec_from_elem!(i64, |x| x == 0);
1432+
impl_spec_from_elem!(i128, |x| x == 0);
1433+
impl_spec_from_elem!(isize, |x| x == 0);
1434+
1435+
impl_spec_from_elem!(u16, |x| x == 0);
1436+
impl_spec_from_elem!(u32, |x| x == 0);
1437+
impl_spec_from_elem!(u64, |x| x == 0);
1438+
impl_spec_from_elem!(u128, |x| x == 0);
1439+
impl_spec_from_elem!(usize, |x| x == 0);
1440+
1441+
impl_spec_from_elem!(f32, |x: f32| x == 0. && x.is_sign_positive());
1442+
impl_spec_from_elem!(f64, |x: f64| x == 0. && x.is_sign_positive());
1443+
13781444
////////////////////////////////////////////////////////////////////////////////
13791445
// Common trait implementations for Vec
13801446
////////////////////////////////////////////////////////////////////////////////

‎src/test/run-pass/auxiliary/allocator-dummy.rs

+5
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,11 @@ pub extern fn __rust_allocate(size: usize, align: usize) -> *mut u8 {
2727
}
2828
}
2929

30+
#[no_mangle]
31+
pub extern fn __rust_allocate_zeroed(size: usize, _align: usize) -> *mut u8 {
32+
unsafe { libc::calloc(size as libc::size_t, 1) as *mut u8 }
33+
}
34+
3035
#[no_mangle]
3136
pub extern fn __rust_deallocate(ptr: *mut u8, old_size: usize, align: usize) {
3237
unsafe {

0 commit comments

Comments
 (0)