library/std/src/sys/pal/windows/alloc.rs - third_party/rust - Git at Google

 use core::mem::MaybeUninit;

 use crate::alloc::{GlobalAlloc, Layout, System};
 use crate::ffi::c_void;
 use crate::ptr;
 use crate::sync::atomic::{AtomicPtr, Ordering};
 use crate::sys::c::{self, windows_targets};
 use crate::sys::common::alloc::{realloc_fallback, MIN_ALIGN};

 #[cfg(test)]
 mod tests;

 // Heap memory management on Windows is done by using the system Heap API (heapapi.h)
 // See https://docs.microsoft.com/windows/win32/api/heapapi/

 // Flag to indicate that the memory returned by `HeapAlloc` should be zeroed.
 const HEAP_ZERO_MEMORY: u32 = 0x00000008;

 // Get a handle to the default heap of the current process, or null if the operation fails.
 //
 // SAFETY: Successful calls to this function within the same process are assumed to
 // always return the same handle, which remains valid for the entire lifetime of the process.
 //
 // See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-getprocessheap
 windows_targets::link!("kernel32.dll" "system" fn GetProcessHeap() -> c::HANDLE);

 // Allocate a block of `dwBytes` bytes of memory from a given heap `hHeap`.
 // The allocated memory may be uninitialized, or zeroed if `dwFlags` is
 // set to `HEAP_ZERO_MEMORY`.
 //
 // Returns a pointer to the newly-allocated memory or null if the operation fails.
 // The returned pointer will be aligned to at least `MIN_ALIGN`.
 //
 // SAFETY:
 //  - `hHeap` must be a non-null handle returned by `GetProcessHeap`.
 //  - `dwFlags` must be set to either zero or `HEAP_ZERO_MEMORY`.
 //
 // Note that `dwBytes` is allowed to be zero, contrary to some other allocators.
 //
 // See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-heapalloc
 windows_targets::link!("kernel32.dll" "system" fn HeapAlloc(hheap: c::HANDLE, dwflags: u32, dwbytes: usize) -> *mut c_void);

 // Reallocate a block of memory behind a given pointer `lpMem` from a given heap `hHeap`,
 // to a block of at least `dwBytes` bytes, either shrinking the block in place,
 // or allocating at a new location, copying memory, and freeing the original location.
 //
 // Returns a pointer to the reallocated memory or null if the operation fails.
 // The returned pointer will be aligned to at least `MIN_ALIGN`.
 // If the operation fails the given block will never have been freed.
 //
 // SAFETY:
 //  - `hHeap` must be a non-null handle returned by `GetProcessHeap`.
 //  - `dwFlags` must be set to zero.
 //  - `lpMem` must be a non-null pointer to an allocated block returned by `HeapAlloc` or
 //     `HeapReAlloc`, that has not already been freed.
 // If the block was successfully reallocated at a new location, pointers pointing to
 // the freed memory, such as `lpMem`, must not be dereferenced ever again.
 //
 // Note that `dwBytes` is allowed to be zero, contrary to some other allocators.
 //
 // See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-heaprealloc
 windows_targets::link!("kernel32.dll" "system" fn HeapReAlloc(
     hheap: c::HANDLE,
     dwflags : u32,
     lpmem: *const c_void,
     dwbytes: usize
 ) -> *mut c_void);

 // Free a block of memory behind a given pointer `lpMem` from a given heap `hHeap`.
 // Returns a nonzero value if the operation is successful, and zero if the operation fails.
 //
 // SAFETY:
 //  - `hHeap` must be a non-null handle returned by `GetProcessHeap`.
 //  - `dwFlags` must be set to zero.
 //  - `lpMem` must be a pointer to an allocated block returned by `HeapAlloc` or `HeapReAlloc`,
 //     that has not already been freed.
 // If the block was successfully freed, pointers pointing to the freed memory, such as `lpMem`,
 // must not be dereferenced ever again.
 //
 // Note that `lpMem` is allowed to be null, which will not cause the operation to fail.
 //
 // See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-heapfree
 windows_targets::link!("kernel32.dll" "system" fn HeapFree(hheap: c::HANDLE, dwflags: u32, lpmem: *const c_void) -> c::BOOL);

 // Cached handle to the default heap of the current process.
 // Either a non-null handle returned by `GetProcessHeap`, or null when not yet initialized or `GetProcessHeap` failed.
 static HEAP: AtomicPtr<c_void> = AtomicPtr::new(ptr::null_mut());

 // Get a handle to the default heap of the current process, or null if the operation fails.
 // If this operation is successful, `HEAP` will be successfully initialized and contain
 // a non-null handle returned by `GetProcessHeap`.
 #[inline]
 fn init_or_get_process_heap() -> c::HANDLE {
     // `HEAP` has not yet been successfully initialized
     let heap = unsafe { GetProcessHeap() };
     if !heap.is_null() {
         // SAFETY: No locking is needed because within the same process,
         // successful calls to `GetProcessHeap` will always return the same value, even on different threads.
         HEAP.store(heap, Ordering::Release);

         // SAFETY: `HEAP` contains a non-null handle returned by `GetProcessHeap`
         heap
     } else {
         // Could not get the current process heap.
         ptr::null_mut()
     }
 }

 /// This is outlined from `process_heap_alloc` so that `process_heap_alloc`
 /// does not need any stack allocations.
 #[inline(never)]
 #[cold]
 extern "C" fn process_heap_init_and_alloc(
     _heap: MaybeUninit<c::HANDLE>, // We pass this argument to match the ABI of `HeapAlloc`
     flags: u32,
     dwBytes: usize,
 ) -> *mut c_void {
     let heap = init_or_get_process_heap();
     if core::intrinsics::unlikely(heap.is_null()) {
         return ptr::null_mut();
     }
     // SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`.
     unsafe { HeapAlloc(heap, flags, dwBytes) }
 }

 #[inline(never)]
 fn process_heap_alloc(
     _heap: MaybeUninit<c::HANDLE>, // We pass this argument to match the ABI of `HeapAlloc`,
     flags: u32,
     dwBytes: usize,
 ) -> *mut c_void {
     let heap = HEAP.load(Ordering::Relaxed);
     if core::intrinsics::likely(!heap.is_null()) {
         // SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`.
         unsafe { HeapAlloc(heap, flags, dwBytes) }
     } else {
         process_heap_init_and_alloc(MaybeUninit::uninit(), flags, dwBytes)
     }
 }

 // Get a non-null handle to the default heap of the current process.
 // SAFETY: `HEAP` must have been successfully initialized.
 #[inline]
 unsafe fn get_process_heap() -> c::HANDLE {
     HEAP.load(Ordering::Acquire)
 }

 // Header containing a pointer to the start of an allocated block.
 // SAFETY: Size and alignment must be <= `MIN_ALIGN`.
 #[repr(C)]
 struct Header(*mut u8);

 // Allocate a block of optionally zeroed memory for a given `layout`.
 // SAFETY: Returns a pointer satisfying the guarantees of `System` about allocated pointers,
 // or null if the operation fails. If this returns non-null `HEAP` will have been successfully
 // initialized.
 #[inline]
 unsafe fn allocate(layout: Layout, zeroed: bool) -> *mut u8 {
     // Allocated memory will be either zeroed or uninitialized.
     let flags = if zeroed { HEAP_ZERO_MEMORY } else { 0 };

     if layout.align() <= MIN_ALIGN {
         // The returned pointer points to the start of an allocated block.
         process_heap_alloc(MaybeUninit::uninit(), flags, layout.size()) as *mut u8
     } else {
         // Allocate extra padding in order to be able to satisfy the alignment.
         let total = layout.align() + layout.size();

         let ptr = process_heap_alloc(MaybeUninit::uninit(), flags, total) as *mut u8;
         if ptr.is_null() {
             // Allocation has failed.
             return ptr::null_mut();
         }

         // Create a correctly aligned pointer offset from the start of the allocated block,
         // and write a header before it.

         let offset = layout.align() - (ptr.addr() & (layout.align() - 1));
         // SAFETY: `MIN_ALIGN` <= `offset` <= `layout.align()` and the size of the allocated
         // block is `layout.align() + layout.size()`. `aligned` will thus be a correctly aligned
         // pointer inside the allocated block with at least `layout.size()` bytes after it and at
         // least `MIN_ALIGN` bytes of padding before it.
         let aligned = unsafe { ptr.add(offset) };
         // SAFETY: Because the size and alignment of a header is <= `MIN_ALIGN` and `aligned`
         // is aligned to at least `MIN_ALIGN` and has at least `MIN_ALIGN` bytes of padding before
         // it, it is safe to write a header directly before it.
         unsafe { ptr::write((aligned as *mut Header).sub(1), Header(ptr)) };

         // SAFETY: The returned pointer does not point to the start of an allocated block,
         // but there is a header readable directly before it containing the location of the start
         // of the block.
         aligned
     }
 }

 // All pointers returned by this allocator have, in addition to the guarantees of `GlobalAlloc`, the
 // following properties:
 //
 // If the pointer was allocated or reallocated with a `layout` specifying an alignment <= `MIN_ALIGN`
 // the pointer will be aligned to at least `MIN_ALIGN` and point to the start of the allocated block.
 //
 // If the pointer was allocated or reallocated with a `layout` specifying an alignment > `MIN_ALIGN`
 // the pointer will be aligned to the specified alignment and not point to the start of the allocated block.
 // Instead there will be a header readable directly before the returned pointer, containing the actual
 // location of the start of the block.
 #[stable(feature = "alloc_system_type", since = "1.28.0")]
 unsafe impl GlobalAlloc for System {
     #[inline]
     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
         // SAFETY: Pointers returned by `allocate` satisfy the guarantees of `System`
         let zeroed = false;
         unsafe { allocate(layout, zeroed) }
     }

     #[inline]
     unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
         // SAFETY: Pointers returned by `allocate` satisfy the guarantees of `System`
         let zeroed = true;
         unsafe { allocate(layout, zeroed) }
     }

     #[inline]
     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
         let block = {
             if layout.align() <= MIN_ALIGN {
                 ptr
             } else {
                 // The location of the start of the block is stored in the padding before `ptr`.

                 // SAFETY: Because of the contract of `System`, `ptr` is guaranteed to be non-null
                 // and have a header readable directly before it.
                 unsafe { ptr::read((ptr as *mut Header).sub(1)).0 }
             }
         };

         // SAFETY: because `ptr` has been successfully allocated with this allocator,
         // `HEAP` must have been successfully initialized.
         let heap = unsafe { get_process_heap() };

         // SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`,
         // `block` is a pointer to the start of an allocated block.
         unsafe { HeapFree(heap, 0, block.cast::<c_void>()) };
     }

     #[inline]
     unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
         if layout.align() <= MIN_ALIGN {
             // SAFETY: because `ptr` has been successfully allocated with this allocator,
             // `HEAP` must have been successfully initialized.
             let heap = unsafe { get_process_heap() };

             // SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`,
             // `ptr` is a pointer to the start of an allocated block.
             // The returned pointer points to the start of an allocated block.
             unsafe { HeapReAlloc(heap, 0, ptr.cast::<c_void>(), new_size).cast::<u8>() }
         } else {
             // SAFETY: `realloc_fallback` is implemented using `dealloc` and `alloc`, which will
             // correctly handle `ptr` and return a pointer satisfying the guarantees of `System`
             unsafe { realloc_fallback(self, ptr, layout, new_size) }
         }
     }
 }
	use core::mem::MaybeUninit;

	use crate::alloc::{GlobalAlloc, Layout, System};
	use crate::ffi::c_void;
	use crate::ptr;
	use crate::sync::atomic::{AtomicPtr, Ordering};
	use crate::sys::c::{self, windows_targets};
	use crate::sys::common::alloc::{realloc_fallback, MIN_ALIGN};

	#[cfg(test)]
	mod tests;

	// Heap memory management on Windows is done by using the system Heap API (heapapi.h)
	// See https://docs.microsoft.com/windows/win32/api/heapapi/

	// Flag to indicate that the memory returned by `HeapAlloc` should be zeroed.
	const HEAP_ZERO_MEMORY: u32 = 0x00000008;

	// Get a handle to the default heap of the current process, or null if the operation fails.
	//
	// SAFETY: Successful calls to this function within the same process are assumed to
	// always return the same handle, which remains valid for the entire lifetime of the process.
	//
	// See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-getprocessheap
	windows_targets::link!("kernel32.dll" "system" fn GetProcessHeap() -> c::HANDLE);

	// Allocate a block of `dwBytes` bytes of memory from a given heap `hHeap`.
	// The allocated memory may be uninitialized, or zeroed if `dwFlags` is
	// set to `HEAP_ZERO_MEMORY`.
	//
	// Returns a pointer to the newly-allocated memory or null if the operation fails.
	// The returned pointer will be aligned to at least `MIN_ALIGN`.
	//
	// SAFETY:
	// - `hHeap` must be a non-null handle returned by `GetProcessHeap`.
	// - `dwFlags` must be set to either zero or `HEAP_ZERO_MEMORY`.
	//
	// Note that `dwBytes` is allowed to be zero, contrary to some other allocators.
	//
	// See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-heapalloc
	windows_targets::link!("kernel32.dll" "system" fn HeapAlloc(hheap: c::HANDLE, dwflags: u32, dwbytes: usize) -> *mut c_void);

	// Reallocate a block of memory behind a given pointer `lpMem` from a given heap `hHeap`,
	// to a block of at least `dwBytes` bytes, either shrinking the block in place,
	// or allocating at a new location, copying memory, and freeing the original location.
	//
	// Returns a pointer to the reallocated memory or null if the operation fails.
	// The returned pointer will be aligned to at least `MIN_ALIGN`.
	// If the operation fails the given block will never have been freed.
	//
	// SAFETY:
	// - `hHeap` must be a non-null handle returned by `GetProcessHeap`.
	// - `dwFlags` must be set to zero.
	// - `lpMem` must be a non-null pointer to an allocated block returned by `HeapAlloc` or
	// `HeapReAlloc`, that has not already been freed.
	// If the block was successfully reallocated at a new location, pointers pointing to
	// the freed memory, such as `lpMem`, must not be dereferenced ever again.
	//
	// Note that `dwBytes` is allowed to be zero, contrary to some other allocators.
	//
	// See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-heaprealloc
	windows_targets::link!("kernel32.dll" "system" fn HeapReAlloc(
	hheap: c::HANDLE,
	dwflags : u32,
	lpmem: *const c_void,
	dwbytes: usize
	) -> *mut c_void);

	// Free a block of memory behind a given pointer `lpMem` from a given heap `hHeap`.
	// Returns a nonzero value if the operation is successful, and zero if the operation fails.
	//
	// SAFETY:
	// - `hHeap` must be a non-null handle returned by `GetProcessHeap`.
	// - `dwFlags` must be set to zero.
	// - `lpMem` must be a pointer to an allocated block returned by `HeapAlloc` or `HeapReAlloc`,
	// that has not already been freed.
	// If the block was successfully freed, pointers pointing to the freed memory, such as `lpMem`,
	// must not be dereferenced ever again.
	//
	// Note that `lpMem` is allowed to be null, which will not cause the operation to fail.
	//
	// See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-heapfree
	windows_targets::link!("kernel32.dll" "system" fn HeapFree(hheap: c::HANDLE, dwflags: u32, lpmem: *const c_void) -> c::BOOL);

	// Cached handle to the default heap of the current process.
	// Either a non-null handle returned by `GetProcessHeap`, or null when not yet initialized or `GetProcessHeap` failed.
	static HEAP: AtomicPtr<c_void> = AtomicPtr::new(ptr::null_mut());

	// Get a handle to the default heap of the current process, or null if the operation fails.
	// If this operation is successful, `HEAP` will be successfully initialized and contain
	// a non-null handle returned by `GetProcessHeap`.
	#[inline]
	fn init_or_get_process_heap() -> c::HANDLE {
	// `HEAP` has not yet been successfully initialized
	let heap = unsafe { GetProcessHeap() };
	if !heap.is_null() {
	// SAFETY: No locking is needed because within the same process,
	// successful calls to `GetProcessHeap` will always return the same value, even on different threads.
	HEAP.store(heap, Ordering::Release);

	// SAFETY: `HEAP` contains a non-null handle returned by `GetProcessHeap`
	heap
	} else {
	// Could not get the current process heap.
	ptr::null_mut()
	}
	}

	/// This is outlined from `process_heap_alloc` so that `process_heap_alloc`
	/// does not need any stack allocations.
	#[inline(never)]
	#[cold]
	extern "C" fn process_heap_init_and_alloc(
	_heap: MaybeUninit<c::HANDLE>, // We pass this argument to match the ABI of `HeapAlloc`
	flags: u32,
	dwBytes: usize,
	) -> *mut c_void {
	let heap = init_or_get_process_heap();
	if core::intrinsics::unlikely(heap.is_null()) {
	return ptr::null_mut();
	}
	// SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`.
	unsafe { HeapAlloc(heap, flags, dwBytes) }
	}

	#[inline(never)]
	fn process_heap_alloc(
	_heap: MaybeUninit<c::HANDLE>, // We pass this argument to match the ABI of `HeapAlloc`,
	flags: u32,
	dwBytes: usize,
	) -> *mut c_void {
	let heap = HEAP.load(Ordering::Relaxed);
	if core::intrinsics::likely(!heap.is_null()) {
	// SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`.
	unsafe { HeapAlloc(heap, flags, dwBytes) }
	} else {
	process_heap_init_and_alloc(MaybeUninit::uninit(), flags, dwBytes)
	}
	}

	// Get a non-null handle to the default heap of the current process.
	// SAFETY: `HEAP` must have been successfully initialized.
	#[inline]
	unsafe fn get_process_heap() -> c::HANDLE {
	HEAP.load(Ordering::Acquire)
	}

	// Header containing a pointer to the start of an allocated block.
	// SAFETY: Size and alignment must be <= `MIN_ALIGN`.
	#[repr(C)]
	struct Header(*mut u8);

	// Allocate a block of optionally zeroed memory for a given `layout`.
	// SAFETY: Returns a pointer satisfying the guarantees of `System` about allocated pointers,
	// or null if the operation fails. If this returns non-null `HEAP` will have been successfully
	// initialized.
	#[inline]
	unsafe fn allocate(layout: Layout, zeroed: bool) -> *mut u8 {
	// Allocated memory will be either zeroed or uninitialized.
	let flags = if zeroed { HEAP_ZERO_MEMORY } else { 0 };

	if layout.align() <= MIN_ALIGN {
	// The returned pointer points to the start of an allocated block.
	process_heap_alloc(MaybeUninit::uninit(), flags, layout.size()) as *mut u8
	} else {
	// Allocate extra padding in order to be able to satisfy the alignment.
	let total = layout.align() + layout.size();

	let ptr = process_heap_alloc(MaybeUninit::uninit(), flags, total) as *mut u8;
	if ptr.is_null() {
	// Allocation has failed.
	return ptr::null_mut();
	}

	// Create a correctly aligned pointer offset from the start of the allocated block,
	// and write a header before it.

	let offset = layout.align() - (ptr.addr() & (layout.align() - 1));
	// SAFETY: `MIN_ALIGN` <= `offset` <= `layout.align()` and the size of the allocated
	// block is `layout.align() + layout.size()`. `aligned` will thus be a correctly aligned
	// pointer inside the allocated block with at least `layout.size()` bytes after it and at
	// least `MIN_ALIGN` bytes of padding before it.
	let aligned = unsafe { ptr.add(offset) };
	// SAFETY: Because the size and alignment of a header is <= `MIN_ALIGN` and `aligned`
	// is aligned to at least `MIN_ALIGN` and has at least `MIN_ALIGN` bytes of padding before
	// it, it is safe to write a header directly before it.
	unsafe { ptr::write((aligned as *mut Header).sub(1), Header(ptr)) };

	// SAFETY: The returned pointer does not point to the start of an allocated block,
	// but there is a header readable directly before it containing the location of the start
	// of the block.
	aligned
	}
	}

	// All pointers returned by this allocator have, in addition to the guarantees of `GlobalAlloc`, the
	// following properties:
	//
	// If the pointer was allocated or reallocated with a `layout` specifying an alignment <= `MIN_ALIGN`
	// the pointer will be aligned to at least `MIN_ALIGN` and point to the start of the allocated block.
	//
	// If the pointer was allocated or reallocated with a `layout` specifying an alignment > `MIN_ALIGN`
	// the pointer will be aligned to the specified alignment and not point to the start of the allocated block.
	// Instead there will be a header readable directly before the returned pointer, containing the actual
	// location of the start of the block.
	#[stable(feature = "alloc_system_type", since = "1.28.0")]
	unsafe impl GlobalAlloc for System {
	#[inline]
	unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
	// SAFETY: Pointers returned by `allocate` satisfy the guarantees of `System`
	let zeroed = false;
	unsafe { allocate(layout, zeroed) }
	}

	#[inline]
	unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
	// SAFETY: Pointers returned by `allocate` satisfy the guarantees of `System`
	let zeroed = true;
	unsafe { allocate(layout, zeroed) }
	}

	#[inline]
	unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
	let block = {
	if layout.align() <= MIN_ALIGN {
	ptr
	} else {
	// The location of the start of the block is stored in the padding before `ptr`.

	// SAFETY: Because of the contract of `System`, `ptr` is guaranteed to be non-null
	// and have a header readable directly before it.
	unsafe { ptr::read((ptr as *mut Header).sub(1)).0 }
	}
	};

	// SAFETY: because `ptr` has been successfully allocated with this allocator,
	// `HEAP` must have been successfully initialized.
	let heap = unsafe { get_process_heap() };

	// SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`,
	// `block` is a pointer to the start of an allocated block.
	unsafe { HeapFree(heap, 0, block.cast::<c_void>()) };
	}

	#[inline]
	unsafe fn realloc(&self, ptr: mut u8, layout: Layout, new_size: usize) -> mut u8 {
	if layout.align() <= MIN_ALIGN {
	// SAFETY: because `ptr` has been successfully allocated with this allocator,
	// `HEAP` must have been successfully initialized.
	let heap = unsafe { get_process_heap() };

	// SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`,
	// `ptr` is a pointer to the start of an allocated block.
	// The returned pointer points to the start of an allocated block.
	unsafe { HeapReAlloc(heap, 0, ptr.cast::<c_void>(), new_size).cast::<u8>() }
	} else {
	// SAFETY: `realloc_fallback` is implemented using `dealloc` and `alloc`, which will
	// correctly handle `ptr` and return a pointer satisfying the guarantees of `System`
	unsafe { realloc_fallback(self, ptr, layout, new_size) }
	}
	}
	}