libavf: add AVirtualMachineRawConfig_setVCpuCount

Bug: 381281047
Test: m
Change-Id: I59356bf13507808e783362d79e7b07559f195ba2
diff --git a/libs/libavf/include/android/virtualization.h b/libs/libavf/include/android/virtualization.h
index 7f07667..2e3abff 100644
--- a/libs/libavf/include/android/virtualization.h
+++ b/libs/libavf/include/android/virtualization.h
@@ -147,6 +147,15 @@
                                             int32_t swiotlbMiB) __INTRODUCED_IN(36);
 
 /**
+ * Set vCPU count. The default is 1.
+ *
+ * \param config a virtual machine config object.
+ * \param n number of vCPUs. Must be positive.
+ */
+void AVirtualMachineRawConfig_setVCpuCount(AVirtualMachineRawConfig* _Nonnull config, int32_t n)
+        __INTRODUCED_IN(36);
+
+/**
  * Set whether the virtual machine's memory will be protected from the host, so the host can't
  * access its memory.
  *
diff --git a/libs/libavf/libavf.map.txt b/libs/libavf/libavf.map.txt
index 2935234..6289b8d 100644
--- a/libs/libavf/libavf.map.txt
+++ b/libs/libavf/libavf.map.txt
@@ -9,6 +9,7 @@
     AVirtualMachineRawConfig_addDisk; # apex llndk
     AVirtualMachineRawConfig_setMemoryMiB; # apex llndk
     AVirtualMachineRawConfig_setSwiotlbMiB; # apex llndk
+    AVirtualMachineRawConfig_setVCpuCount; # apex llndk
     AVirtualMachineRawConfig_setProtectedVm; # apex llndk
     AVirtualMachineRawConfig_setHypervisorSpecificAuthMethod; # apex llndk
     AVirtualMachineRawConfig_addCustomMemoryBackingFile; # apex llndk
diff --git a/libs/libavf/src/lib.rs b/libs/libavf/src/lib.rs
index 8d60f3c..af891e8 100644
--- a/libs/libavf/src/lib.rs
+++ b/libs/libavf/src/lib.rs
@@ -23,8 +23,8 @@
 
 use android_system_virtualizationservice::{
     aidl::android::system::virtualizationservice::{
-        DiskImage::DiskImage, IVirtualizationService::IVirtualizationService,
-        VirtualMachineConfig::VirtualMachineConfig,
+        CpuTopology::CpuTopology, DiskImage::DiskImage,
+        IVirtualizationService::IVirtualizationService, VirtualMachineConfig::VirtualMachineConfig,
         VirtualMachineRawConfig::VirtualMachineRawConfig,
     },
     binder::{ParcelFileDescriptor, Strong},
@@ -205,6 +205,22 @@
     config.swiotlbMib = swiotlb_mib;
 }
 
+/// Set vCPU count.
+///
+/// # Safety
+/// `config` must be a pointer returned by `AVirtualMachineRawConfig_create`.
+#[no_mangle]
+pub unsafe extern "C" fn AVirtualMachineRawConfig_setVCpuCount(
+    config: *mut VirtualMachineRawConfig,
+    n: i32,
+) {
+    // SAFETY: `config` is assumed to be a valid, non-null pointer returned by
+    // AVirtualMachineRawConfig_create. It's the only reference to the object.
+    let config = unsafe { &mut *config };
+    config.cpuTopology = CpuTopology::CUSTOM;
+    config.customVcpuCount = n;
+}
+
 /// Set whether a virtual machine is protected or not.
 ///
 /// # Safety