From 1668abb4008cb869df2b669e52ac8847762db07a Mon Sep 17 00:00:00 2001
From: John Schock <joschock@microsoft.com>
Date: Fri, 20 Oct 2023 16:37:29 -0700
Subject: [PATCH] Update HelloWorldRustDxe to permit compiling for unit tests

---
 MsCorePkg/HelloWorldRustDxe/Cargo.toml  |  1 -
 MsCorePkg/HelloWorldRustDxe/src/main.rs | 65 ++++++++++++++++---------
 2 files changed, 41 insertions(+), 25 deletions(-)

diff --git a/MsCorePkg/HelloWorldRustDxe/Cargo.toml b/MsCorePkg/HelloWorldRustDxe/Cargo.toml
index 005c87f6ee..501c791ec3 100644
--- a/MsCorePkg/HelloWorldRustDxe/Cargo.toml
+++ b/MsCorePkg/HelloWorldRustDxe/Cargo.toml
@@ -6,7 +6,6 @@ edition = "2021"
 [[bin]]
 name = "HelloWorldRustDxe"
 path = "src/main.rs"
-test = false
 
 [dependencies]
 r-efi = {workspace=true}
diff --git a/MsCorePkg/HelloWorldRustDxe/src/main.rs b/MsCorePkg/HelloWorldRustDxe/src/main.rs
index 4cff82374c..fd35fd021a 100644
--- a/MsCorePkg/HelloWorldRustDxe/src/main.rs
+++ b/MsCorePkg/HelloWorldRustDxe/src/main.rs
@@ -8,39 +8,56 @@
 //!
 //! SPDX-License-Identifier: BSD-2-Clause-Patent
 //!
-#![no_std]
-#![no_main]
+#![cfg_attr(target_os = "uefi", no_std)]
+#![cfg_attr(target_os = "uefi", no_main)]
 #![allow(non_snake_case)]
 
-extern crate alloc;
+#[cfg(target_os = "uefi")]
+mod uefi_entry {
+  extern crate alloc;
 
-use alloc::vec;
-use core::panic::PanicInfo;
-use r_efi::efi::Status;
-use rust_advanced_logger_dxe::{debugln, init_debug, DEBUG_INFO};
+  use alloc::vec;
+  use core::panic::PanicInfo;
+  use r_efi::efi::Status;
+  use rust_advanced_logger_dxe::{debugln, init_debug, DEBUG_INFO};
 
-#[no_mangle]
-pub extern "efiapi" fn efi_main(
-  _image_handle: *const core::ffi::c_void,
-  _system_table: *const r_efi::system::SystemTable,
-) -> u64 {
-  rust_boot_services_allocator_dxe::GLOBAL_ALLOCATOR.init(unsafe { (*_system_table).boot_services });
-  init_debug(unsafe { (*_system_table).boot_services });
+  #[no_mangle]
+  pub extern "efiapi" fn efi_main(
+    _image_handle: *const core::ffi::c_void,
+    _system_table: *const r_efi::system::SystemTable,
+  ) -> u64 {
+    rust_boot_services_allocator_dxe::GLOBAL_ALLOCATOR.init(unsafe { (*_system_table).boot_services });
+    init_debug(unsafe { (*_system_table).boot_services });
 
-  debugln!(DEBUG_INFO, "Hello, World. This is Rust in UEFI.");
+    debugln!(DEBUG_INFO, "Hello, World. This is Rust in UEFI.");
 
-  debugln!(DEBUG_INFO, "file: {:} line: {:} as hex: {:x}", file!(), line!(), line!());
+    debugln!(DEBUG_INFO, "file: {:} line: {:} as hex: {:x}", file!(), line!(), line!());
 
-  let mut foo = vec!["asdf", "xyzpdq", "abcdefg", "thxyzb"];
+    let mut foo = vec!["asdf", "xyzpdq", "abcdefg", "thxyzb"];
 
-  debugln!(DEBUG_INFO, "Unsorted vec: {:?}", foo);
-  foo.sort();
-  debugln!(DEBUG_INFO, "Sorted vec: {:?}", foo);
+    debugln!(DEBUG_INFO, "Unsorted vec: {:?}", foo);
+    foo.sort();
+    debugln!(DEBUG_INFO, "Sorted vec: {:?}", foo);
 
-  Status::SUCCESS.as_usize() as u64
+    Status::SUCCESS.as_usize() as u64
+  }
+
+  #[panic_handler]
+  fn panic(_info: &PanicInfo) -> ! {
+    loop {}
+  }
+}
+
+#[cfg(not(target_os = "uefi"))]
+fn main() {
+  //do nothing.
 }
 
-#[panic_handler]
-fn panic(_info: &PanicInfo) -> ! {
-  loop {}
+#[cfg(test)]
+mod test {
+
+  #[test]
+  fn sample_test() {
+    assert_eq!(1 + 1, 2);
+  }
 }