From c07671edb4a8734fac8c7515965a03070e6faee1 Mon Sep 17 00:00:00 2001
From: Wei Liu <liuwe@microsoft.com>
Date: Tue, 7 May 2024 00:51:22 +0000
Subject: [PATCH] hypervisor: kvm: introduce a check_extension macro

That reduces code repetition.

Signed-off-by: Wei Liu <liuwe@microsoft.com>
---
 hypervisor/src/kvm/aarch64/mod.rs | 10 ++++++++--
 hypervisor/src/kvm/x86_64/mod.rs  | 30 +++++++++++++-----------------
 2 files changed, 21 insertions(+), 19 deletions(-)

diff --git a/hypervisor/src/kvm/aarch64/mod.rs b/hypervisor/src/kvm/aarch64/mod.rs
index fb5021a62..61ff79443 100644
--- a/hypervisor/src/kvm/aarch64/mod.rs
+++ b/hypervisor/src/kvm/aarch64/mod.rs
@@ -106,9 +106,15 @@ pub fn is_system_register(regid: u64) -> bool {
 }
 
 pub fn check_required_kvm_extensions(kvm: &Kvm) -> KvmResult<()> {
-    if !kvm.check_extension(Cap::OneReg) {
-        return Err(KvmError::CapabilityMissing(Cap::OneReg));
+    macro_rules! check_extension {
+        ($cap:expr) => {
+            if !kvm.check_extension($cap) {
+                return Err(KvmError::CapabilityMissing($cap));
+            }
+        };
     }
+
+    check_extension!(Cap::OneReg);
     Ok(())
 }
 
diff --git a/hypervisor/src/kvm/x86_64/mod.rs b/hypervisor/src/kvm/x86_64/mod.rs
index 8c019b3ab..0d539551c 100644
--- a/hypervisor/src/kvm/x86_64/mod.rs
+++ b/hypervisor/src/kvm/x86_64/mod.rs
@@ -32,24 +32,20 @@ pub use {
 /// Check KVM extension for Linux
 ///
 pub fn check_required_kvm_extensions(kvm: &Kvm) -> KvmResult<()> {
-    if !kvm.check_extension(Cap::GetTscKhz) {
-        return Err(KvmError::CapabilityMissing(Cap::GetTscKhz));
-    }
-    if !kvm.check_extension(Cap::ImmediateExit) {
-        return Err(KvmError::CapabilityMissing(Cap::ImmediateExit));
-    }
-    if !kvm.check_extension(Cap::SetIdentityMapAddr) {
-        return Err(KvmError::CapabilityMissing(Cap::SetIdentityMapAddr));
-    }
-    if !kvm.check_extension(Cap::SetTssAddr) {
-        return Err(KvmError::CapabilityMissing(Cap::SetTssAddr));
-    }
-    if !kvm.check_extension(Cap::SplitIrqchip) {
-        return Err(KvmError::CapabilityMissing(Cap::SplitIrqchip));
-    }
-    if !kvm.check_extension(Cap::TscDeadlineTimer) {
-        return Err(KvmError::CapabilityMissing(Cap::TscDeadlineTimer));
+    macro_rules! check_extension {
+        ($cap:expr) => {
+            if !kvm.check_extension($cap) {
+                return Err(KvmError::CapabilityMissing($cap));
+            }
+        };
     }
+
+    check_extension!(Cap::GetTscKhz);
+    check_extension!(Cap::ImmediateExit);
+    check_extension!(Cap::SetIdentityMapAddr);
+    check_extension!(Cap::SetTssAddr);
+    check_extension!(Cap::SplitIrqchip);
+    check_extension!(Cap::TscDeadlineTimer);
     Ok(())
 }
 #[derive(Clone, Serialize, Deserialize)]