refactor(aarch64/vcpu): set one register at a time

Replace `set_registers` with `set_register` in order
to have more control over each register restoration
process.

Signed-off-by: Egor Lazarchuk <yegorlz@amazon.co.uk>

Author: ShadowCurse

src/vmm/src/arch/aarch64/vcpu.rs

@@ -179,17 +179,15 @@ pub fn get_all_registers_ids(vcpufd: &VcpuFd) -> Result<Vec<u64>, VcpuError> {
     }
 }
 
-/// Set the state of the system registers.
+/// Set the state of one system register.
 ///
 /// # Arguments
 ///
-/// * `regs` - Slice of registers to be set.
-pub fn set_registers(vcpufd: &VcpuFd, regs: &Aarch64RegisterVec) -> Result<(), VcpuError> {
-    for reg in regs.iter() {
-        vcpufd
-            .set_one_reg(reg.id, reg.as_slice())
-            .map_err(|e| VcpuError::SetOneReg(reg.id, e))?;
-    }
+/// * `reg` - Register to be set.
+pub fn set_register(vcpufd: &VcpuFd, reg: Aarch64RegisterRef) -> Result<(), VcpuError> {
+    vcpufd
+        .set_one_reg(reg.id, reg.as_slice())
+        .map_err(|e| VcpuError::SetOneReg(reg.id, e))?;
     Ok(())
 }
 
@@ -277,7 +275,9 @@ mod tests {
 
         vcpu.vcpu_init(&kvi).unwrap();
         get_all_registers(&vcpu, &mut regs).unwrap();
-        set_registers(&vcpu, &regs).unwrap();
+        for reg in regs.iter() {
+            set_register(&vcpu, reg).unwrap();
+        }
     }
 
     #[test]

src/vmm/src/vstate/vcpu/aarch64.rs

@@ -16,7 +16,7 @@ use crate::arch::aarch64::regs::{
 };
 use crate::arch::aarch64::vcpu::{
     get_all_registers, get_all_registers_ids, get_mpidr, get_mpstate, get_registers, set_mpstate,
-    set_registers, setup_boot_regs, VcpuError as ArchError,
+    set_register, setup_boot_regs, VcpuError as ArchError,
 };
 use crate::cpu_config::aarch64::custom_cpu_template::VcpuFeatures;
 use crate::cpu_config::templates::CpuConfiguration;
@@ -191,9 +191,11 @@ impl KvmVcpu {
         };
 
         self.init_vcpu_fd(&kvi)?;
-
         self.kvi = state.kvi;
-        set_registers(&self.fd, &state.regs).map_err(KvmVcpuError::RestoreState)?;
+
+        for reg in state.regs.iter() {
+            set_register(&self.fd, reg).map_err(KvmVcpuError::RestoreState)?;
+        }
         set_mpstate(&self.fd, state.mp_state).map_err(KvmVcpuError::RestoreState)?;
         Ok(())
     }