Skip to content

Commit e518d03

Browse files
authored
Avoid double allocation when passing strings via IntoParam (#1713)
* Use an iterator instead of double allocation for strings * Add a missing trailing new line * Use `T::default` instead of `Default::default` for clarity * Rename `from_iter` to `string_from_iter` for clarity * Run `cargo fmt` * Make string_from_iter from memory safe * Encoder is fused because Chain is fused * Remove unsafe from string_from_iter as it is memory-safe * Incorporate ryancerium and rylev's suggestions * Correct documentation error
1 parent 9deec0e commit e518d03

File tree

3 files changed

+26
-14
lines changed

3 files changed

+26
-14
lines changed

crates/libs/windows/src/core/heap.rs

Lines changed: 23 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -32,20 +32,32 @@ pub unsafe fn heap_free(ptr: RawPtr) {
3232
}
3333
}
3434

35-
/// Copy a slice of `T` into a freshly allocated buffer with an additional default `T` at the end.
35+
/// Copy len elements of an iterator of type `T` into a freshly allocated buffer.
3636
///
37-
/// Returns a pointer to the beginning of the buffer
37+
/// Returns a pointer to the beginning of the buffer. This pointer must be freed when done using `heap_free`.
3838
///
3939
/// # Panics
4040
///
41-
/// This function panics if the heap allocation fails or if the pointer returned from
42-
/// the heap allocation is not properly aligned to `T`.
43-
pub fn heap_string<T: Copy + Default + Sized>(slice: &[T]) -> *const T {
44-
unsafe {
45-
let buffer = heap_alloc((slice.len() + 1) * std::mem::size_of::<T>()).expect("could not allocate string") as *mut T;
46-
assert!(buffer.align_offset(std::mem::align_of::<T>()) == 0, "heap allocated buffer is not properly aligned");
47-
buffer.copy_from_nonoverlapping(slice.as_ptr(), slice.len());
48-
buffer.add(slice.len()).write(T::default());
49-
buffer
41+
/// This function panics if the heap allocation fails, the alignment requirements of 'T' surpass
42+
/// 8 (HeapAlloc's alignment).
43+
pub fn alloc_from_iter<I, T>(iter: I, len: usize) -> *const T
44+
where
45+
I: Iterator<Item = T>,
46+
T: Copy,
47+
{
48+
// alignment of memory returned by HeapAlloc is at least 8
49+
// Source: https://docs.microsoft.com/en-us/windows/win32/api/heapapi/nf-heapapi-heapalloc
50+
// Ensure that T has sufficient alignment requirements
51+
assert!(std::mem::align_of::<T>() <= 8, "T alignment surpasses HeapAlloc alignment");
52+
53+
let ptr = heap_alloc(len * std::mem::size_of::<T>()).expect("could not allocate string") as *mut T;
54+
55+
for (offset, c) in iter.take(len).enumerate() {
56+
// SAFETY: ptr points to an allocation object of size `len`, indices accessed are always lower than `len`
57+
unsafe {
58+
ptr.add(offset).write(c);
59+
}
5060
}
61+
62+
ptr
5163
}

crates/libs/windows/src/core/pcstr.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ unsafe impl Abi for PCSTR {
4545
#[cfg(feature = "alloc")]
4646
impl<'a> IntoParam<'a, PCSTR> for &str {
4747
fn into_param(self) -> Param<'a, PCSTR> {
48-
Param::Boxed(PCSTR(heap_string(self.as_bytes())))
48+
Param::Boxed(PCSTR(alloc_from_iter(self.as_bytes().iter().copied().chain(core::iter::once(0)), self.len() + 1)))
4949
}
5050
}
5151
#[cfg(feature = "alloc")]

crates/libs/windows/src/core/pcwstr.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ unsafe impl Abi for PCWSTR {
4545
#[cfg(feature = "alloc")]
4646
impl<'a> IntoParam<'a, PCWSTR> for &str {
4747
fn into_param(self) -> Param<'a, PCWSTR> {
48-
Param::Boxed(PCWSTR(heap_string(&self.encode_utf16().collect::<alloc::vec::Vec<u16>>())))
48+
Param::Boxed(PCWSTR(alloc_from_iter(self.encode_utf16().chain(core::iter::once(0)), self.len() + 1)))
4949
}
5050
}
5151
#[cfg(feature = "alloc")]
@@ -58,7 +58,7 @@ impl<'a> IntoParam<'a, PCWSTR> for alloc::string::String {
5858
impl<'a> IntoParam<'a, PCWSTR> for &::std::ffi::OsStr {
5959
fn into_param(self) -> Param<'a, PCWSTR> {
6060
use ::std::os::windows::ffi::OsStrExt;
61-
Param::Boxed(PCWSTR(heap_string(&self.encode_wide().collect::<alloc::vec::Vec<u16>>())))
61+
Param::Boxed(PCWSTR(alloc_from_iter(self.encode_wide().chain(core::iter::once(0)), self.len() + 1)))
6262
}
6363
}
6464
#[cfg(feature = "alloc")]

0 commit comments

Comments
 (0)