From aa12a34823d8017154466883d479274130ad1f69 Mon Sep 17 00:00:00 2001
From: Tony Arcieri <bascule@gmail.com>
Date: Fri, 6 Sep 2024 07:37:11 -0600
Subject: [PATCH] aarch64-dit: have `set_dit_enabled` return previous state
 (#1109)

This is more inlining-friendly
---
 aarch64-dit/src/lib.rs | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/aarch64-dit/src/lib.rs b/aarch64-dit/src/lib.rs
index 8b5a4b66..0eef23ba 100644
--- a/aarch64-dit/src/lib.rs
+++ b/aarch64-dit/src/lib.rs
@@ -44,9 +44,7 @@ impl Dit {
     #[must_use]
     pub fn enable(&self) -> Guard<'_> {
         let was_enabled = if self.is_supported() {
-            let was_enabled = unsafe { get_dit_enabled() };
-            unsafe { set_dit_enabled() };
-            was_enabled
+            unsafe { set_dit_enabled() }
         } else {
             false
         };
@@ -102,9 +100,13 @@ unsafe fn get_dit_enabled() -> bool {
 }
 
 /// Enable DIT for the current thread.
+///
+/// Returns the previous DIT state prior to enabling DIT.
 #[target_feature(enable = "dit")]
-unsafe fn set_dit_enabled() {
+unsafe fn set_dit_enabled() -> bool {
+    let was_enabled = get_dit_enabled();
     asm!("msr DIT, #1", options(nomem, nostack, preserves_flags));
+    was_enabled
 }
 
 /// Restore DIT state depending on the enabled bit.
@@ -153,7 +155,8 @@ mod tests {
         let dit_enabled = unsafe { get_dit_enabled() };
         assert!(!dit_enabled);
 
-        unsafe { set_dit_enabled() };
+        let was_enabled = unsafe { set_dit_enabled() };
+        assert!(!was_enabled);
         let dit_enabled = unsafe { get_dit_enabled() };
         assert!(dit_enabled);