From 2d8635f04a7c6283420d59157c19209a9aba5569 Mon Sep 17 00:00:00 2001
From: Michael Zhao <michael.zhao@arm.com>
Date: Tue, 21 Jun 2022 11:32:28 +0800
Subject: [PATCH] hypervisor: Refactor `system_registers` on AArch64

Function `system_registers` took mutable vector reference and modified
the vector content. Now change the definition to `get/set` style.
And rename to `get/set_sys_regs` to align with other functions.

Signed-off-by: Michael Zhao <michael.zhao@arm.com>
---
 hypervisor/src/cpu.rs     |  4 ++--
 hypervisor/src/kvm/mod.rs | 11 ++++++-----
 vmm/src/cpu.rs            | 13 ++++++++-----
 3 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/hypervisor/src/cpu.rs b/hypervisor/src/cpu.rs
index 68e3fba71..e9b659ebe 100644
--- a/hypervisor/src/cpu.rs
+++ b/hypervisor/src/cpu.rs
@@ -429,12 +429,12 @@ pub trait Vcpu: Send + Sync {
     /// Save the state of the system registers.
     ///
     #[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
-    fn system_registers(&self, state: &mut Vec<Register>) -> Result<()>;
+    fn get_sys_regs(&self) -> Result<Vec<Register>>;
     ///
     /// Restore the state of the system registers.
     ///
     #[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
-    fn set_system_registers(&self, state: &[Register]) -> Result<()>;
+    fn set_sys_regs(&self, state: &[Register]) -> Result<()>;
     ///
     /// Read the MPIDR - Multiprocessor Affinity Register.
     ///
diff --git a/hypervisor/src/kvm/mod.rs b/hypervisor/src/kvm/mod.rs
index 0ce5b1190..1458651e7 100644
--- a/hypervisor/src/kvm/mod.rs
+++ b/hypervisor/src/kvm/mod.rs
@@ -1573,9 +1573,10 @@ impl cpu::Vcpu for KvmVcpu {
     /// Save the state of the system registers.
     ///
     #[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
-    fn system_registers(&self, state: &mut Vec<Register>) -> cpu::Result<()> {
+    fn get_sys_regs(&self) -> cpu::Result<Vec<Register>> {
         // Call KVM_GET_REG_LIST to get all registers available to the guest. For ArmV8 there are
         // around 500 registers.
+        let mut state: Vec<Register> = Vec::new();
         let mut reg_list = RegList::new(500).unwrap();
         self.fd
             .get_reg_list(&mut reg_list)
@@ -1601,13 +1602,13 @@ impl cpu::Vcpu for KvmVcpu {
             });
         }
 
-        Ok(())
+        Ok(state)
     }
     ///
     /// Restore the state of the system registers.
     ///
     #[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
-    fn set_system_registers(&self, state: &[Register]) -> cpu::Result<()> {
+    fn set_sys_regs(&self, state: &[Register]) -> cpu::Result<()> {
         for reg in state {
             self.fd
                 .set_one_reg(reg.id, reg.addr)
@@ -1811,7 +1812,7 @@ impl cpu::Vcpu for KvmVcpu {
             ..Default::default()
         };
         state.core_regs = self.get_regs()?;
-        self.system_registers(&mut state.sys_regs)?;
+        state.sys_regs = self.get_sys_regs()?;
 
         Ok(state)
     }
@@ -1905,7 +1906,7 @@ impl cpu::Vcpu for KvmVcpu {
     #[cfg(target_arch = "aarch64")]
     fn set_state(&self, state: &CpuState) -> cpu::Result<()> {
         self.set_regs(&state.core_regs)?;
-        self.set_system_registers(&state.sys_regs)?;
+        self.set_sys_regs(&state.sys_regs)?;
         self.set_mp_state(state.mp_state)?;
 
         Ok(())
diff --git a/vmm/src/cpu.rs b/vmm/src/cpu.rs
index cc5b12d94..76e6cdd53 100644
--- a/vmm/src/cpu.rs
+++ b/vmm/src/cpu.rs
@@ -2537,10 +2537,10 @@ mod tests {
 
         // Must fail when vcpu is not initialized yet.
         let mut state: Vec<kvm_one_reg> = Vec::new();
-        let res = vcpu.system_registers(&mut state);
+        let res = vcpu.get_sys_regs();
         assert!(res.is_err());
         assert_eq!(
-            format!("{}", res.unwrap_err()),
+            format!("{}", res.as_ref().unwrap_err()),
             "Failed to retrieve list of registers: Exec format error (os error 8)"
         );
 
@@ -2548,7 +2548,7 @@ mod tests {
             id: MPIDR_EL1,
             addr: 0x00,
         });
-        let res = vcpu.set_system_registers(&state);
+        let res = vcpu.set_sys_regs(&state);
         assert!(res.is_err());
         assert_eq!(
             format!("{}", res.unwrap_err()),
@@ -2556,14 +2556,17 @@ mod tests {
         );
 
         vcpu.vcpu_init(&kvi).unwrap();
-        assert!(vcpu.system_registers(&mut state).is_ok());
+        let res = vcpu.get_sys_regs();
+        assert!(res.is_ok());
+        state = res.unwrap();
+
         let initial_mpidr: u64 = vcpu.read_mpidr().expect("Fail to read mpidr");
         assert!(state.contains(&kvm_one_reg {
             id: MPIDR_EL1,
             addr: initial_mpidr
         }));
 
-        assert!(vcpu.set_system_registers(&state).is_ok());
+        assert!(vcpu.set_sys_regs(&state).is_ok());
         let mpidr: u64 = vcpu.read_mpidr().expect("Fail to read mpidr");
         assert_eq!(initial_mpidr, mpidr);
     }