Add support for building `std::os::uefi` docs

Signed-off-by: Ayush Singh <ayushsingh1325@gmail.com>

6 files changed, 33 insertions, 28 deletions

diff --git a/library/std/src/os/mod.rs b/library/std/src/os/mod.rs
index 11ad21515fd..d1a73503899 100644
--- a/library/std/src/os/mod.rs
+++ b/library/std/src/os/mod.rs
@@ -143,6 +143,7 @@ pub mod solid;
 #[path = "ios/mod.rs"]
 pub(crate) mod tvos;
 #[cfg(target_os = "uefi")]
+#[cfg(any(target_os = "uefi", doc))]
 pub mod uefi;
 #[cfg(target_os = "vita")]
 pub mod vita;
diff --git a/library/std/src/os/uefi/env.rs b/library/std/src/os/uefi/env.rs
index e62441596f3..41dc7ccdb6a 100644
--- a/library/std/src/os/uefi/env.rs
+++ b/library/std/src/os/uefi/env.rs
@@ -32,41 +32,45 @@ pub unsafe fn init_globals(handle: NonNull<c_void>, system_table: NonNull<c_void
 }
 
 /// Get the SystemTable Pointer.
+/// If you want to use `BootServices` then please use [`boot_services`] as it performs some
+/// additional checks.
+///
 /// Note: This function panics if the System Table or Image Handle is not initialized
 pub fn system_table() -> NonNull<c_void> {
     try_system_table().unwrap()
 }
 
 /// Get the ImageHandle Pointer.
+///
 /// Note: This function panics if the System Table or Image Handle is not initialized
 pub fn image_handle() -> NonNull<c_void> {
     try_image_handle().unwrap()
 }
 
-/// Get the SystemTable Pointer.
-/// This function is mostly intended for places where panic is not an option
-pub(crate) fn try_system_table() -> Option<NonNull<crate::ffi::c_void>> {
-    GLOBALS.get().map(|x| x.0)
-}
-
-/// Get the SystemHandle Pointer.
-/// This function is mostly intended for places where panic is not an option
-pub(crate) fn try_image_handle() -> Option<NonNull<crate::ffi::c_void>> {
-    GLOBALS.get().map(|x| x.1)
-}
-
 /// Get the BootServices Pointer.
 /// This function also checks if `ExitBootServices` has already been called.
-pub(crate) fn boot_services() -> Option<NonNull<r_efi::efi::BootServices>> {
+pub fn boot_services() -> Option<NonNull<c_void>> {
     if BOOT_SERVICES_FLAG.get() {
         let system_table: NonNull<r_efi::efi::SystemTable> = try_system_table()?.cast();
         let boot_services = unsafe { (*system_table.as_ptr()).boot_services };
-        NonNull::new(boot_services)
+        NonNull::new(boot_services).map(|x| x.cast())
     } else {
         None
     }
 }
 
+/// Get the SystemTable Pointer.
+/// This function is mostly intended for places where panic is not an option
+pub(crate) fn try_system_table() -> Option<NonNull<c_void>> {
+    GLOBALS.get().map(|x| x.0)
+}
+
+/// Get the SystemHandle Pointer.
+/// This function is mostly intended for places where panic is not an option
+pub(crate) fn try_image_handle() -> Option<NonNull<c_void>> {
+    GLOBALS.get().map(|x| x.1)
+}
+
 pub(crate) fn enable_boot_services() {
     BOOT_SERVICES_FLAG.set(true);
 }
diff --git a/library/std/src/os/uefi/mod.rs b/library/std/src/os/uefi/mod.rs
index fc0468f25fc..8ef05eee1f4 100644
--- a/library/std/src/os/uefi/mod.rs
+++ b/library/std/src/os/uefi/mod.rs
@@ -1,6 +1,7 @@
 //! Platform-specific extensions to `std` for UEFI.
 
 #![unstable(feature = "uefi_std", issue = "100499")]
+#![doc(cfg(target_os = "uefi"))]
 
 pub mod env;
 #[path = "../windows/ffi.rs"]
diff --git a/library/std/src/sys/uefi/alloc.rs b/library/std/src/sys/uefi/alloc.rs
index d6eb371033c..789e3cbd81a 100644
--- a/library/std/src/sys/uefi/alloc.rs
+++ b/library/std/src/sys/uefi/alloc.rs
@@ -1,7 +1,7 @@
 //! Global Allocator for UEFI.
 //! Uses [r-efi-alloc](https://crates.io/crates/r-efi-alloc)
 
-use crate::alloc::{handle_alloc_error, GlobalAlloc, Layout, System};
+use crate::alloc::{GlobalAlloc, Layout, System};
 
 const MEMORY_TYPE: u32 = r_efi::efi::LOADER_DATA;
 
@@ -13,11 +13,8 @@ unsafe impl GlobalAlloc for System {
             return crate::ptr::null_mut();
         }
 
-        let system_table = match crate::os::uefi::env::try_system_table() {
-            None => return crate::ptr::null_mut(),
-            Some(x) => x.as_ptr() as *mut _,
-        };
-
+        // If boot services is valid then SystemTable is not null.
+        let system_table = crate::os::uefi::env::system_table().as_ptr().cast();
         // The caller must ensure non-0 layout
         unsafe { r_efi_alloc::raw::alloc(system_table, layout, MEMORY_TYPE) }
     }
@@ -28,10 +25,8 @@ unsafe impl GlobalAlloc for System {
             return;
         }
 
-        let system_table = match crate::os::uefi::env::try_system_table() {
-            None => handle_alloc_error(layout),
-            Some(x) => x.as_ptr() as *mut _,
-        };
+        // If boot services is valid then SystemTable is not null.
+        let system_table = crate::os::uefi::env::system_table().as_ptr().cast();
         // The caller must ensure non-0 layout
         unsafe { r_efi_alloc::raw::dealloc(system_table, ptr, layout) }
     }
diff --git a/library/std/src/sys/uefi/helpers.rs b/library/std/src/sys/uefi/helpers.rs
index 28a82394226..2108484a40c 100644
--- a/library/std/src/sys/uefi/helpers.rs
+++ b/library/std/src/sys/uefi/helpers.rs
@@ -46,7 +46,7 @@ pub(crate) fn locate_handles(mut guid: Guid) -> io::Result<Vec<NonNull<crate::ff
         if r.is_error() { Err(status_to_io_error(r)) } else { Ok(()) }
     }
 
-    let boot_services = boot_services().ok_or(BOOT_SERVICES_UNAVAILABLE)?;
+    let boot_services = boot_services().ok_or(BOOT_SERVICES_UNAVAILABLE)?.cast();
     let mut buf_len = 0usize;
 
     // This should always fail since the size of buffer is 0. This call should update the buf_len
@@ -82,7 +82,8 @@ pub(crate) fn open_protocol<T>(
     handle: NonNull<crate::ffi::c_void>,
     mut protocol_guid: Guid,
 ) -> io::Result<NonNull<T>> {
-    let boot_services = boot_services().ok_or(BOOT_SERVICES_UNAVAILABLE)?;
+    let boot_services: NonNull<efi::BootServices> =
+        boot_services().ok_or(BOOT_SERVICES_UNAVAILABLE)?.cast();
     let system_handle = uefi::env::image_handle();
     let mut protocol: MaybeUninit<*mut T> = MaybeUninit::uninit();
 
@@ -279,7 +280,8 @@ pub(crate) fn create_event(
     handler: Option<efi::EventNotify>,
     context: *mut crate::ffi::c_void,
 ) -> io::Result<NonNull<crate::ffi::c_void>> {
-    let boot_services = boot_services().ok_or(BOOT_SERVICES_UNAVAILABLE)?;
+    let boot_services: NonNull<efi::BootServices> =
+        boot_services().ok_or(BOOT_SERVICES_UNAVAILABLE)?.cast();
     let mut exit_boot_service_event: r_efi::efi::Event = crate::ptr::null_mut();
     let r = unsafe {
         let create_event = (*boot_services.as_ptr()).create_event;
@@ -294,7 +296,8 @@ pub(crate) fn create_event(
 }
 
 pub(crate) fn close_event(evt: NonNull<crate::ffi::c_void>) -> io::Result<()> {
-    let boot_services = boot_services().ok_or(BOOT_SERVICES_UNAVAILABLE)?;
+    let boot_services: NonNull<efi::BootServices> =
+        boot_services().ok_or(BOOT_SERVICES_UNAVAILABLE)?.cast();
     let r = unsafe {
         let close_event = (*boot_services.as_ptr()).close_event;
         (close_event)(evt.as_ptr())
diff --git a/library/std/src/sys/uefi/mod.rs b/library/std/src/sys/uefi/mod.rs
index fa33b681a19..6fd2c5f21cb 100644
--- a/library/std/src/sys/uefi/mod.rs
+++ b/library/std/src/sys/uefi/mod.rs
@@ -129,6 +129,7 @@ pub fn abort_internal() -> ! {
     if let (Some(boot_services), Some(handle)) =
         (uefi::env::boot_services(), uefi::env::try_image_handle())
     {
+        let boot_services: NonNull<r_efi::efi::BootServices> = boot_services.cast();
         let _ = unsafe {
             ((*boot_services.as_ptr()).exit)(
                 handle.as_ptr(),