Implement Instant for UEFI

- Uses Timestamp Protocol if present. Else use rdtsc for x86 and x86-64

Signed-off-by: Ayush Singh <[email protected]>

Author: Ayush1325

library/std/src/lib.rs

@@ -263,6 +263,10 @@
 #![cfg_attr(any(windows, target_os = "uefi"), feature(round_char_boundary))]
 #![cfg_attr(target_os = "xous", feature(slice_ptr_len))]
 #![cfg_attr(target_family = "wasm", feature(stdarch_wasm_atomic_wait))]
+#![cfg_attr(
+    all(any(target_arch = "x86_64", target_arch = "x86"), target_os = "uefi"),
+    feature(stdarch_x86_has_cpuid)
+)]
 //
 // Language features:
 // tidy-alphabetical-start

library/std/src/sys/pal/uefi/time.rs

@@ -14,6 +14,15 @@ pub const UNIX_EPOCH: SystemTime = SystemTime(Duration::from_secs(0));
 
 impl Instant {
     pub fn now() -> Instant {
+        // If we have a timestamp protocol, use it.
+        if let Some(x) = instant_internal::timestamp_protocol() {
+            return x;
+        }
+
+        if let Some(x) = instant_internal::platform_specific() {
+            return x;
+        }
+
         panic!("time not implemented on this platform")
     }
 
@@ -103,3 +112,111 @@ pub(crate) mod system_time_internal {
         Duration::new(utc_epoch, t.nanosecond)
     }
 }
+
+pub(crate) mod instant_internal {
+    use super::super::helpers;
+    use super::*;
+    use crate::mem::MaybeUninit;
+    use crate::ptr::NonNull;
+    use crate::sync::atomic::{AtomicPtr, Ordering};
+    use crate::sync::OnceLock;
+    use crate::sys_common::mul_div_u64;
+    use r_efi::protocols::timestamp;
+
+    const NS_PER_SEC: u64 = 1_000_000_000;
+
+    pub fn timestamp_protocol() -> Option<Instant> {
+        fn try_handle(handle: NonNull<crate::ffi::c_void>) -> Option<u64> {
+            let protocol: NonNull<timestamp::Protocol> =
+                helpers::open_protocol(handle, timestamp::PROTOCOL_GUID).ok()?;
+            let mut properties: MaybeUninit<timestamp::Properties> = MaybeUninit::uninit();
+
+            let r = unsafe { ((*protocol.as_ptr()).get_properties)(properties.as_mut_ptr()) };
+            if r.is_error() {
+                return None;
+            }
+
+            let freq = unsafe { properties.assume_init().frequency };
+            let ts = unsafe { ((*protocol.as_ptr()).get_timestamp)() };
+            Some(mul_div_u64(ts, NS_PER_SEC, freq))
+        }
+
+        static LAST_VALID_HANDLE: AtomicPtr<crate::ffi::c_void> =
+            AtomicPtr::new(crate::ptr::null_mut());
+
+        if let Some(handle) = NonNull::new(LAST_VALID_HANDLE.load(Ordering::Acquire)) {
+            if let Some(ns) = try_handle(handle) {
+                return Some(Instant(Duration::from_nanos(ns)));
+            }
+        }
+
+        if let Ok(handles) = helpers::locate_handles(timestamp::PROTOCOL_GUID) {
+            for handle in handles {
+                if let Some(ns) = try_handle(handle) {
+                    LAST_VALID_HANDLE.store(handle.as_ptr(), Ordering::Release);
+                    return Some(Instant(Duration::from_nanos(ns)));
+                }
+            }
+        }
+
+        None
+    }
+
+    pub fn platform_specific() -> Option<Instant> {
+        if cfg!(target_arch = "x86_64") {
+            timestamp_rdtsc().map(Instant)
+        } else if cfg!(target_arch = "x86") {
+            timestamp_rdtsc().map(Instant)
+        } else {
+            None
+        }
+    }
+
+    #[cfg(target_arch = "x86_64")]
+    fn timestamp_rdtsc() -> Option<Duration> {
+        if !crate::arch::x86_64::has_cpuid() {
+            return None;
+        }
+
+        static FREQUENCY: OnceLock<u64> = OnceLock::new();
+
+        // Get Frequency in Mhz
+        // Inspired by [`edk2/UefiCpuPkg/Library/CpuTimerLib/CpuTimerLib.c`](https://github.com/tianocore/edk2/blob/master/UefiCpuPkg/Library/CpuTimerLib/CpuTimerLib.c)
+        let freq = FREQUENCY
+            .get_or_try_init(|| {
+                let cpuid = unsafe { crate::arch::x86_64::__cpuid(0x15) };
+                if cpuid.eax == 0 || cpuid.ebx == 0 || cpuid.ecx == 0 {
+                    return Err(());
+                }
+                Ok(mul_div_u64(cpuid.ecx as u64, cpuid.ebx as u64, cpuid.eax as u64))
+            })
+            .ok()?;
+
+        let ts = unsafe { crate::arch::x86_64::_rdtsc() };
+        let ns = mul_div_u64(ts, 1000, *freq);
+        Some(Duration::from_nanos(ns))
+    }
+
+    #[cfg(target_arch = "x86")]
+    fn timestamp_rdtsc() -> Option<Duration> {
+        if !crate::arch::x86::has_cpuid() {
+            return None;
+        }
+
+        static FREQUENCY: OnceLock<u64> = OnceLock::new();
+
+        let freq = FREQUENCY
+            .get_or_try_init(|| {
+                let cpuid = unsafe { crate::arch::x86::__cpuid(0x15) };
+                if cpuid.eax == 0 || cpuid.ebx == 0 || cpuid.ecx == 0 {
+                    return Err(());
+                }
+                Ok(mul_div_u64(cpuid.ecx as u64, cpuid.ebx as u64, cpuid.eax as u64))
+            })
+            .ok()?;
+
+        let ts = unsafe { crate::arch::x86::_rdtsc() };
+        let ns = mul_div_u64(ts, 1000, *freq);
+        Some(Duration::from_nanos(ns))
+    }
+}