From c01f71de21486ff8e558fa905cc1aa42264ac6a5 Mon Sep 17 00:00:00 2001
From: Putta Khunchalee <ultimaweapon@outlook.com>
Date: Sat, 7 Sep 2024 01:06:36 +0700
Subject: [PATCH] Refactors CpuStates::set_pstate

---
 src/core/Cargo.toml              |  1 +
 src/core/src/vmm/aarch64.rs      | 11 +++++++++--
 src/core/src/vmm/hv/aarch64.rs   | 16 ++++++++++++++++
 src/core/src/vmm/hv/macos/cpu.rs | 12 ++----------
 src/core/src/vmm/hv/mod.rs       |  4 +---
 5 files changed, 29 insertions(+), 15 deletions(-)

diff --git a/src/core/Cargo.toml b/src/core/Cargo.toml
index 8d9ab6d5..dd4912cd 100644
--- a/src/core/Cargo.toml
+++ b/src/core/Cargo.toml
@@ -7,6 +7,7 @@ edition = "2021"
 crate-type = ["staticlib"]
 
 [dependencies]
+bitfield-struct = "0.8.0"
 humansize = "2.1.3"
 libc = "0.2.155"
 obconf = { path = "../obconf" }
diff --git a/src/core/src/vmm/aarch64.rs b/src/core/src/vmm/aarch64.rs
index b085befe..02ebca9a 100644
--- a/src/core/src/vmm/aarch64.rs
+++ b/src/core/src/vmm/aarch64.rs
@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT OR Apache-2.0
-use super::hv::{Cpu, CpuFeats, CpuStates};
+use super::hv::{Cpu, CpuFeats, CpuStates, Pstate};
 use super::hw::RamMap;
 use super::MainCpuError;
 
@@ -31,7 +31,14 @@ pub fn setup_main_cpu(
     // Set PSTATE so the PE run in AArch64 mode. Not sure why we need M here since the document said
     // it is ignore. See https://gist.github.com/imbushuo/51b09e61ecd7b7ac063853ad65cedf34 where
     // M = 5 came from.
-    states.set_pstate(true, true, true, true, 0b101);
+    states.set_pstate(
+        Pstate::new()
+            .with_m(0b101)
+            .with_f(true)
+            .with_i(true)
+            .with_a(true)
+            .with_d(true),
+    );
 
     // Enable MMU to enable virtual address and set TCR_EL1.
     states.set_sctlr_el1(true);
diff --git a/src/core/src/vmm/hv/aarch64.rs b/src/core/src/vmm/hv/aarch64.rs
index c6cb493f..dc2d9a32 100644
--- a/src/core/src/vmm/hv/aarch64.rs
+++ b/src/core/src/vmm/hv/aarch64.rs
@@ -1,4 +1,5 @@
 // SPDX-License-Identifier: MIT OR Apache-2.0
+use bitfield_struct::bitfield;
 
 /// Features available on a PE.
 pub struct CpuFeats {
@@ -37,3 +38,18 @@ pub struct CpuFeats {
     /// - `0b0101`: 48 bits, 256TB.
     pub pa_range: u8,
 }
+
+/// Represents a value of `PSTATE`.
+#[bitfield(u32)]
+pub struct Pstate {
+    #[bits(4)]
+    pub m: u8,
+    #[bits(2)]
+    __: u8,
+    pub f: bool,
+    pub i: bool,
+    pub a: bool,
+    pub d: bool,
+    #[bits(22)]
+    __: u32,
+}
diff --git a/src/core/src/vmm/hv/macos/cpu.rs b/src/core/src/vmm/hv/macos/cpu.rs
index 95297e6c..8954287a 100644
--- a/src/core/src/vmm/hv/macos/cpu.rs
+++ b/src/core/src/vmm/hv/macos/cpu.rs
@@ -302,16 +302,8 @@ impl<'a, 'b> CpuStates for HfStates<'a, 'b> {
     }
 
     #[cfg(target_arch = "aarch64")]
-    fn set_pstate(&mut self, d: bool, a: bool, i: bool, f: bool, m: u8) {
-        let d: u64 = d.into();
-        let a: u64 = a.into();
-        let i: u64 = i.into();
-        let f: u64 = f.into();
-        let m: u64 = m.into();
-
-        assert_eq!(m & 0b11110000, 0);
-
-        self.pstate = State::Dirty(d << 9 | a << 8 | i << 7 | f << 6 | m);
+    fn set_pstate(&mut self, v: crate::vmm::hv::Pstate) {
+        self.pstate = State::Dirty(v.into_bits().into());
     }
 
     #[cfg(target_arch = "aarch64")]
diff --git a/src/core/src/vmm/hv/mod.rs b/src/core/src/vmm/hv/mod.rs
index dbb69960..2c96a6fb 100644
--- a/src/core/src/vmm/hv/mod.rs
+++ b/src/core/src/vmm/hv/mod.rs
@@ -104,10 +104,8 @@ pub trait CpuStates {
     #[cfg(target_arch = "x86_64")]
     fn set_ss(&mut self, p: bool);
 
-    /// # Panics
-    /// If `m` larger than 4 bits.
     #[cfg(target_arch = "aarch64")]
-    fn set_pstate(&mut self, d: bool, a: bool, i: bool, f: bool, m: u8);
+    fn set_pstate(&mut self, v: Pstate);
 
     #[cfg(target_arch = "aarch64")]
     fn set_sctlr_el1(&mut self, m: bool);