From 578f6fcb222d57babc1a444575dce3a3250e29c8 Mon Sep 17 00:00:00 2001
From: Changyuan Lyu <changyuanl@google.com>
Date: Fri, 18 Apr 2025 00:15:40 -0700
Subject: [PATCH] refactor(virtio): move Register to VirtioPciDevice

Signed-off-by: Changyuan Lyu <changyuanl@google.com>
---
 alioth/src/virtio/dev/dev.rs | 18 +++++++-----------
 alioth/src/virtio/pci.rs     | 13 ++++++++-----
 2 files changed, 15 insertions(+), 16 deletions(-)

diff --git a/alioth/src/virtio/dev/dev.rs b/alioth/src/virtio/dev/dev.rs
index c0d024b7..f48b3858 100644
--- a/alioth/src/virtio/dev/dev.rs
+++ b/alioth/src/virtio/dev/dev.rs
@@ -123,7 +123,7 @@ where
     pub name: Arc<str>,
     pub id: DeviceId,
     pub device_config: Arc<dyn Mmio>,
-    pub reg: Arc<Register>,
+    pub device_feature: u64,
     pub queue_regs: Arc<[Queue]>,
     pub ioeventfds: Arc<[(E, bool)]>,
     pub shared_mem_regions: Option<Arc<MemRegion>>,
@@ -163,16 +163,12 @@ where
         let name = name.into();
         let id = D::DEVICE_ID;
         let device_config = dev.config();
-        let mut dev_feat = dev.feature();
+        let mut device_feature = dev.feature();
         if restricted_memory {
-            dev_feat |= VirtioFeature::ACCESS_PLATFORM.bits()
+            device_feature |= VirtioFeature::ACCESS_PLATFORM.bits()
         } else {
-            dev_feat &= !VirtioFeature::ACCESS_PLATFORM.bits()
+            device_feature &= !VirtioFeature::ACCESS_PLATFORM.bits()
         }
-        let reg = Arc::new(Register {
-            device_feature: dev_feat,
-            ..Default::default()
-        });
         let num_queues = dev.num_queues();
         let queue_regs = (0..num_queues).map(|_| Queue {
             size: AtomicU16::new(QUEUE_SIZE_MAX),
@@ -195,13 +191,13 @@ where
             dev.spawn_worker(event_rx, memory, queue_regs.clone(), ioeventfds.clone())?;
         log::debug!(
             "{name}: created with {:x?} {:x?}",
-            VirtioFeature::from_bits_retain(reg.device_feature & !D::Feature::all().bits()),
-            D::Feature::from_bits_truncate(reg.device_feature)
+            VirtioFeature::from_bits_retain(device_feature & !D::Feature::all().bits()),
+            D::Feature::from_bits_truncate(device_feature)
         );
         let virtio_dev = VirtioDevice {
             name,
             id,
-            reg,
+            device_feature,
             queue_regs,
             ioeventfds,
             worker_handle: Some(handle),
diff --git a/alioth/src/virtio/pci.rs b/alioth/src/virtio/pci.rs
index 8d3a6483..afc94aaf 100644
--- a/alioth/src/virtio/pci.rs
+++ b/alioth/src/virtio/pci.rs
@@ -180,7 +180,7 @@ where
     M: MsiSender,
 {
     name: Arc<str>,
-    reg: Arc<Register>,
+    reg: Register,
     queues: Arc<[Queue]>,
     irq_sender: Arc<PciIrqSender<M>>,
     event_tx: Sender<WakeEvent<PciIrqSender<M>>>,
@@ -244,7 +244,7 @@ where
     }
 
     fn read(&self, offset: u64, size: u8) -> mem::Result<u64> {
-        let reg = &*self.reg;
+        let reg = &self.reg;
         let ret = match (offset as usize, size as usize) {
             VirtioCommonCfg::LAYOUT_DEVICE_FEATURE_SELECT => {
                 reg.device_feature_sel.load(Ordering::Acquire) as u64
@@ -368,7 +368,7 @@ where
     }
 
     fn write(&self, offset: u64, size: u8, val: u64) -> mem::Result<Action> {
-        let reg = &*self.reg;
+        let reg = &self.reg;
         match (offset as usize, size as usize) {
             VirtioCommonCfg::LAYOUT_DEVICE_FEATURE_SELECT => {
                 reg.device_feature_sel.store(val as u8, Ordering::Release);
@@ -812,7 +812,10 @@ where
 
         let registers = Arc::new(VirtioPciRegisterMmio {
             name: dev.name.clone(),
-            reg: dev.reg.clone(),
+            reg: Register {
+                device_feature: dev.device_feature,
+                ..Default::default()
+            },
             event_tx: dev.event_tx.clone(),
             waker: dev.waker.clone(),
             queues: dev.queue_regs.clone(),
@@ -879,7 +882,7 @@ where
     fn reset(&self) -> pci::Result<()> {
         self.registers.wake_up_dev(WakeEvent::Reset);
         self.registers.reset();
-        self.dev.reg.status.store(0, Ordering::Release);
+        self.registers.reg.status.store(0, Ordering::Release);
         Ok(())
     }
 }