From 1ced46da7955468f956f11a58d53c2cc52b9e33a Mon Sep 17 00:00:00 2001
From: li041 <lxh050013@outlook.com>
Date: Thu, 30 Jul 2026 10:59:41 +0800
Subject: [PATCH] zone: track guest-to-physical CPU mappings

---
 src/zone.rs | 45 +++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 43 insertions(+), 2 deletions(-)

diff --git a/src/zone.rs b/src/zone.rs
index b9e68bc87..1e2efefac 100644
--- a/src/zone.rs
+++ b/src/zone.rs
@@ -119,6 +119,10 @@ pub struct ZoneInner {
     mmio: Vec<MMIOConfig>,
     cpu_num: usize,
     cpu_set: CpuSet,
+    // Guest-visible CPU ids are not guaranteed to be dense or equal to pCPU ids.
+    // Keep explicit maps here; CpuSet only records physical CPU ownership.
+    guest_to_phys_cpu: [Option<usize>; MAX_CPU_NUM],
+    phys_to_guest_cpu: [Option<usize>; MAX_CPU_NUM],
     irq_bitmap: [u32; 1024 / 32],
     gpm: MemorySet<Stage2PageTable>,
     iommu_pt: Option<MemorySet<Stage2PageTable>>,
@@ -174,6 +178,8 @@ impl ZoneInner {
             mmio: Vec::new(),
             cpu_num: 0,
             cpu_set: CpuSet::new(MAX_CPU_NUM as usize, 0),
+            guest_to_phys_cpu: [None; MAX_CPU_NUM],
+            phys_to_guest_cpu: [None; MAX_CPU_NUM],
             irq_bitmap: [0; 1024 / 32],
             iommu_pt: if cfg!(iommu) {
                 Some(new_s2_memory_set())
@@ -276,6 +282,36 @@ impl ZoneInner {
         &mut self.cpu_set
     }
 
+    pub fn map_cpu(&mut self, guest_cpu: usize, phys_cpu: usize) {
+        if guest_cpu >= MAX_CPU_NUM || phys_cpu >= MAX_CPU_NUM {
+            warn!(
+                "ignore invalid CPU map guest_cpu={}, phys_cpu={}",
+                guest_cpu, phys_cpu
+            );
+            return;
+        }
+        if let Some(old_phys_cpu) = self.guest_to_phys_cpu[guest_cpu] {
+            if self.phys_to_guest_cpu[old_phys_cpu] == Some(guest_cpu) {
+                self.phys_to_guest_cpu[old_phys_cpu] = None;
+            }
+        }
+        if let Some(old_guest_cpu) = self.phys_to_guest_cpu[phys_cpu] {
+            if self.guest_to_phys_cpu[old_guest_cpu] == Some(phys_cpu) {
+                self.guest_to_phys_cpu[old_guest_cpu] = None;
+            }
+        }
+        self.guest_to_phys_cpu[guest_cpu] = Some(phys_cpu);
+        self.phys_to_guest_cpu[phys_cpu] = Some(guest_cpu);
+    }
+
+    pub fn guest_to_phys_cpu(&self, guest_cpu: usize) -> Option<usize> {
+        self.guest_to_phys_cpu.get(guest_cpu).copied().flatten()
+    }
+
+    pub fn phys_to_guest_cpu(&self, phys_cpu: usize) -> Option<usize> {
+        self.phys_to_guest_cpu.get(phys_cpu).copied().flatten()
+    }
+
     pub fn irq_bitmap(&self) -> &[u32; 1024 / 32] {
         &self.irq_bitmap
     }
@@ -798,7 +834,7 @@ pub fn zone_create(config: &HvZoneConfig) -> HvResult<Arc<Zone>> {
     zone.mmio_init(&config.arch_config);
 
     let mut cpu_num = 0;
-    for cpu_id in config.cpus().iter() {
+    for (guest_cpu, cpu_id) in config.cpus().iter().enumerate() {
         if let Some(existing_zone) = get_cpu_data(*cpu_id as _).zone.clone() {
             return hv_result_err!(
                 EBUSY,
@@ -809,7 +845,12 @@ pub fn zone_create(config: &HvZoneConfig) -> HvResult<Arc<Zone>> {
                 )
             );
         }
-        zone.write().cpu_set_mut().set_bit(*cpu_id as _);
+        let cpu_id = *cpu_id as usize;
+        let mut zone_inner = zone.write();
+        zone_inner.cpu_set_mut().set_bit(cpu_id);
+        // The config lists physical CPUs, while guests use dense CPU ids from
+        // zero. Preserve that distinction for non-root SMP zones.
+        zone_inner.map_cpu(guest_cpu, cpu_id);
         cpu_num += 1;
     }
     zone.write().set_cpu_num(cpu_num);