uefi: Allocator: handle ZST

We used std::alloc::System as template.

[0] https://doc.rust-lang.org/src/std/alloc.rs.html#137

Author: phip1611

Diff for: uefi-test-runner/src/boot/memory.rs

@@ -10,7 +10,11 @@ pub fn test() {
     global::alloc_vec();
     global::alloc_alignment();
 
-    test_memory_map();
+    #[cfg(feature = "unstable")]
+    {
+        allocator_api::alloc_zst();
+        allocator_api::alloc_normal();
+    }
 }
 
 /// Tests that directly use UEFI boot services to allocate memory.
@@ -133,3 +137,38 @@ mod global {
         }
     }
 }
+
+/// Tests the `allocator_api` on the UEFI allocator.
+#[cfg(feature = "unstable")]
+mod allocator_api {
+    use core::alloc::Layout;
+    use uefi::allocator::Allocator;
+
+    pub fn alloc_zst() {
+        let layout = Layout::from_size_align(0, 1024).unwrap();
+        let ptr = <Allocator as core::alloc::Allocator>::allocate(&Allocator, layout).unwrap();
+        assert_eq!(ptr.len(), 0);
+        assert_eq!(ptr.as_ptr().cast::<u8>().align_offset(layout.align()), 0);
+    }
+
+    pub fn alloc_normal() {
+        let layout = Layout::from_size_align(64, 64).unwrap();
+        let allocation =
+            <Allocator as core::alloc::Allocator>::allocate(&Allocator, layout).unwrap();
+        assert_eq!(allocation.len(), 64);
+        assert_eq!(
+            allocation
+                .as_ptr()
+                .cast::<u8>()
+                .align_offset(layout.align()),
+            0
+        );
+
+        unsafe {
+            core::ptr::write_bytes(allocation.as_ptr().cast::<u8>(), 42, allocation.len());
+        }
+        unsafe {
+            assert_eq!(allocation.as_ref()[42], 42);
+        }
+    }
+}

Diff for: uefi-test-runner/src/main.rs

@@ -2,6 +2,7 @@
 
 #![no_std]
 #![no_main]
+#![cfg_attr(feature = "unstable", feature(allocator_api))]
 
 #[macro_use]
 extern crate log;

Diff for: uefi/src/allocator.rs

@@ -173,13 +173,34 @@ unsafe impl GlobalAlloc for Allocator {
 #[cfg(feature = "unstable")]
 unsafe impl core::alloc::Allocator for Allocator {
     fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, core::alloc::AllocError> {
-        let ptr = unsafe { <Allocator as GlobalAlloc>::alloc(self, layout) };
-        NonNull::new(ptr)
-            .ok_or(core::alloc::AllocError)
-            .map(|ptr| NonNull::slice_from_raw_parts(ptr, layout.size()))
+        // Stable alternative for Layout::dangling()
+        fn dangling_for_layout(layout: Layout) -> NonNull<u8> {
+            let align = layout.align();
+            // SAFETY: align is non-zero, so align as usize is a valid address for a NonNull.
+            unsafe {
+                let ptr = align as *mut u8;
+                NonNull::new_unchecked(ptr)
+            }
+        }
+
+        match layout.size() {
+            0 => Ok(NonNull::slice_from_raw_parts(
+                dangling_for_layout(layout),
+                0,
+            )),
+            // SAFETY: `layout` is non-zero in size,
+            size => {
+                let ptr = unsafe { <Allocator as GlobalAlloc>::alloc(self, layout) };
+                NonNull::new(ptr)
+                    .ok_or(core::alloc::AllocError)
+                    .map(|ptr| NonNull::slice_from_raw_parts(ptr, size))
+            }
+        }
     }
 
     unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
-        unsafe { <Allocator as GlobalAlloc>::dealloc(self, ptr.as_ptr(), layout) }
+        if layout.size() != 0 {
+            unsafe { <Allocator as GlobalAlloc>::dealloc(self, ptr.as_ptr(), layout) }
+        }
     }
 }