fix: compilation in aarch64

Changes we did for supporting older snapshot formats, did not really
compile on ARM systems. Fix the compilation issues. The issues were
mainly bad re-exports.

Signed-off-by: Babis Chalios <babis.chalios@e2b.dev>

Author: bchalios

src/vmm/src/arch/aarch64/gic/gicv3/regs/its_regs.rs

@@ -80,7 +80,7 @@ pub fn its_restore_tables(its_fd: &DeviceFd) -> Result<(), GicError> {
 }
 
 /// ITS registers that we save/restore during snapshot
-#[derive(Debug, Default, Serialize, Deserialize)]
+#[derive(Debug, Default, Clone, Serialize, Deserialize)]
 pub struct ItsRegisterState {
     iidr: u64,
     cbaser: u64,

src/vmm/src/arch/aarch64/gic/mod.rs

@@ -8,7 +8,8 @@ mod regs;
 use gicv2::GICv2;
 use gicv3::GICv3;
 use kvm_ioctls::{DeviceFd, VmFd};
-pub use regs::GicState;
+pub use regs::{GicRegState, GicState, GicVcpuState, VgicSysRegsState};
+pub use gicv3::regs::its_regs::ItsRegisterState;
 
 use super::layout;
 

src/vmm/src/arch/aarch64/gic/regs.rs

@@ -12,20 +12,23 @@ use serde::{Deserialize, Serialize};
 use crate::arch::aarch64::gic::GicError;
 use crate::arch::aarch64::gic::gicv3::regs::its_regs::ItsRegisterState;
 
-#[derive(Debug, Serialize, Deserialize)]
+/// Serializable state for a block of GIC registers.
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct GicRegState<T> {
     pub(crate) chunks: Vec<T>,
 }
 
 /// Structure for serializing the state of the Vgic ICC regs
-#[derive(Debug, Default, Serialize, Deserialize)]
+#[derive(Debug, Default, Clone, Serialize, Deserialize)]
 pub struct VgicSysRegsState {
+    /// Main ICC system registers.
     pub main_icc_regs: Vec<GicRegState<u64>>,
+    /// AP ICC system registers (one entry per priority group).
     pub ap_icc_regs: Vec<Option<GicRegState<u64>>>,
 }
 
 /// Structure used for serializing the state of the GIC registers.
-#[derive(Debug, Default, Serialize, Deserialize)]
+#[derive(Debug, Default, Clone, Serialize, Deserialize)]
 pub struct GicState {
     /// The state of the distributor registers.
     pub dist: Vec<GicRegState<u32>>,
@@ -36,9 +39,11 @@ pub struct GicState {
 }
 
 /// Structure used for serializing the state of the GIC registers for a specific vCPU.
-#[derive(Debug, Default, Serialize, Deserialize)]
+#[derive(Debug, Default, Clone, Serialize, Deserialize)]
 pub struct GicVcpuState {
+    /// Redistributor registers for this vCPU.
     pub rdist: Vec<GicRegState<u32>>,
+    /// ICC (CPU interface) system registers for this vCPU.
     pub icc: VgicSysRegsState,
 }
 

src/vmm/src/persist/v1_10/aarch64.rs

@@ -3,20 +3,14 @@
 
 use serde::{Deserialize, Serialize};
 
-use super::{KvmCapability, MMIODeviceInfo};
+use crate::cpu_config::templates::KvmCapability;
+use super::MMIODeviceInfo;
 
 // Types that are identical across all versions — canonical definitions in v1_14.
-pub use crate::v1_14::{
-    StaticCpuTemplate,
-    DeviceType,
-    GicRegState,
-    VgicSysRegsState,
-    GicVcpuState,
-    Aarch64RegisterVec,
-};
+pub use crate::persist::v1_14::DeviceType;
 
 // Types that are identical in v1.10 and v1.12 — canonical definitions in v1_12.
-pub use crate::v1_12::{
+pub use crate::persist::v1_12::{
     // aarch64 GicState is identical in v1.10 and v1.12 (gains its_state in v1.14)
     GicState,
     // aarch64 VcpuState is identical in v1.10 and v1.12 (gains pvtime_ipa in v1.14)

src/vmm/src/persist/v1_12/aarch64.rs

@@ -7,12 +7,11 @@ use serde::{Deserialize, Serialize};
 use super::{GuestMemoryState, MMIODeviceInfo};
 
 // Types that are canonical in v1_14 and unchanged through all versions
-pub use crate::v1_14::{
+pub use crate::persist::v1_14::{
     // Legacy device type enum
     DeviceType,
     // GIC helper types (GicState itself changed — its_state added — so redefined in v1_14)
     GicRegState,
-    VgicSysRegsState,
     GicVcpuState,
     // Register vector with custom serde
     Aarch64RegisterVec,

src/vmm/src/persist/v1_12/mod.rs

@@ -19,6 +19,7 @@
 use serde::{Deserialize, Serialize};
 
 use super::v1_10;
+#[cfg(target_arch = "x86_64")]
 use crate::arch::VcpuState;
 use crate::devices::acpi::vmgenid::VMGenIDState;
 use crate::devices::virtio::balloon::persist::BalloonConfigSpaceState;

src/vmm/src/persist/v1_14/aarch64.rs

@@ -1,18 +1,25 @@
 // Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-use kvm_bindings::{kvm_mp_state, kvm_vcpu_init};
 use serde::{Deserialize, Serialize};
 
-use crate::convert::{ConvertError, irq_to_gsi};
-use crate::v1_12;
+use super::{ACPIDeviceManagerState, ConvertError, GuestMemoryState, MMIODeviceInfo,
+            ResourceAllocator, irq_to_gsi};
+use crate::devices::acpi::vmgenid::VMGenIDState;
+use crate::persist::v1_12;
 
-use super::{
-    ACPIDeviceManagerState, GuestMemoryState, MMIODeviceInfo, ResourceAllocator, VMGenIDState,
-};
+// ───────────────────────────────────────────────────────────────────
+// Re-export runtime types — v1.14 snapshot format matches the runtime format.
+// These are used by v1.12 (and v1.10 via v1.12) as canonical type definitions.
+// ───────────────────────────────────────────────────────────────────
+
+pub use crate::arch::aarch64::gic::{GicRegState, GicState, GicVcpuState};
+pub use crate::arch::aarch64::regs::Aarch64RegisterVec;
+pub use crate::arch::aarch64::vcpu::VcpuState;
+pub use crate::arch::aarch64::vm::VmState;
 
 // ───────────────────────────────────────────────────────────────────
-// StaticCpuTemplate — canonical definition (identical in v1.10, v1.12, v1.14)
+// StaticCpuTemplate — aarch64-specific snapshot enum (same in v1.10, v1.12, v1.14)
 // ───────────────────────────────────────────────────────────────────
 
 #[derive(Debug, Default, Clone, PartialEq, Eq, Serialize, Deserialize)]
@@ -23,7 +30,7 @@ pub enum StaticCpuTemplate {
 }
 
 // ───────────────────────────────────────────────────────────────────
-// aarch64 legacy device types — canonical definitions
+// DeviceType — aarch64 legacy device type enum (snapshot format)
 // ───────────────────────────────────────────────────────────────────
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -33,94 +40,34 @@ pub enum DeviceType {
     Rtc,
 }
 
-// ───────────────────────────────────────────────────────────────────
-// GIC helper types — canonical definitions (unchanged since v1.10)
-// ───────────────────────────────────────────────────────────────────
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-#[serde(bound(serialize = "T: Serialize", deserialize = "T: for<'a> Deserialize<'a>"))]
-pub struct GicRegState<T: Serialize + for<'a> Deserialize<'a>> {
-    pub chunks: Vec<T>,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct VgicSysRegsState {
-    pub main_icc_regs: Vec<GicRegState<u64>>,
-    pub ap_icc_regs: Vec<Option<GicRegState<u64>>>,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct GicVcpuState {
-    pub rdist: Vec<GicRegState<u32>>,
-    pub icc: VgicSysRegsState,
-}
-
-// ───────────────────────────────────────────────────────────────────
-// aarch64 register vector — canonical definition (unchanged since v1.10)
-// ───────────────────────────────────────────────────────────────────
-
-/// aarch64 register vector with custom serde: serialized as (Vec<u64>, Vec<u8>)
-#[derive(Debug, Clone)]
-pub struct Aarch64RegisterVec {
-    pub ids: Vec<u64>,
-    pub data: Vec<u8>,
-}
-
-impl Serialize for Aarch64RegisterVec {
-    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
-        (&self.ids, &self.data).serialize(serializer)
-    }
-}
-
-impl<'de> Deserialize<'de> for Aarch64RegisterVec {
-    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
-        let (ids, data) = <(Vec<u64>, Vec<u8>)>::deserialize(deserializer)?;
-        Ok(Aarch64RegisterVec { ids, data })
+impl From<DeviceType> for crate::arch::DeviceType {
+    fn from(dt: DeviceType) -> Self {
+        match dt {
+            DeviceType::Virtio(n) => crate::arch::DeviceType::Virtio(n),
+            DeviceType::Serial => crate::arch::DeviceType::Serial,
+            DeviceType::Rtc => crate::arch::DeviceType::Rtc,
+        }
     }
 }
 
 // ───────────────────────────────────────────────────────────────────
-// aarch64 ConnectedLegacyState (uses updated MMIODeviceInfo with gsi)
+// ConnectedLegacyState — convert v1.12 snapshot type to runtime type
 // ───────────────────────────────────────────────────────────────────
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct ConnectedLegacyState {
-    pub type_: DeviceType,
-    pub device_info: MMIODeviceInfo,
-}
-
-impl From<v1_12::ConnectedLegacyState> for ConnectedLegacyState {
+impl From<v1_12::ConnectedLegacyState> for crate::device_manager::persist::ConnectedLegacyState {
     fn from(s: v1_12::ConnectedLegacyState) -> Self {
-        ConnectedLegacyState {
-            type_: s.type_,
+        crate::device_manager::persist::ConnectedLegacyState {
+            type_: crate::arch::DeviceType::from(s.type_),
             device_info: MMIODeviceInfo::from(s.device_info),
         }
     }
 }
 
 // ───────────────────────────────────────────────────────────────────
-// aarch64 GIC state (v1.14: adds its_state)
-// GicRegState, VgicSysRegsState, GicVcpuState are defined above
+// GIC state (aarch64, v1.14: adds its_state)
+// GicState is the runtime type (re-exported above); conversion from v1.12 is here.
 // ───────────────────────────────────────────────────────────────────
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct ItsRegisterState {
-    pub iidr: u64,
-    pub cbaser: u64,
-    pub creadr: u64,
-    pub cwriter: u64,
-    pub baser: [u64; 8],
-    pub ctlr: u64,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct GicState {
-    pub dist: Vec<GicRegState<u32>>,
-    pub gic_vcpu_states: Vec<GicVcpuState>,
-    /// ITS state (GICv3 only). None for GICv2 or when converted from v1.12.
-    pub its_state: Option<ItsRegisterState>,
-}
-
 impl GicState {
     pub(crate) fn from(old_state: v1_12::GicState) -> GicState {
         GicState {
@@ -133,17 +80,9 @@ impl GicState {
 
 // ───────────────────────────────────────────────────────────────────
 // vCPU state (aarch64, v1.14: gains pvtime_ipa)
+// VcpuState is the runtime type (re-exported above); conversion from v1.12 is here.
 // ───────────────────────────────────────────────────────────────────
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct VcpuState {
-    pub mp_state: kvm_mp_state,
-    pub regs: Aarch64RegisterVec,
-    pub mpidr: u64,
-    pub kvi: kvm_vcpu_init,
-    pub pvtime_ipa: Option<u64>,
-}
-
 impl VcpuState {
     pub(crate) fn from(old_state: v1_12::VcpuState) -> VcpuState {
         VcpuState {
@@ -157,7 +96,7 @@ impl VcpuState {
 }
 
 // ───────────────────────────────────────────────────────────────────
-// ACPI device state impl (aarch64: no vmclock)
+// ACPI device state (aarch64: no vmclock)
 // ───────────────────────────────────────────────────────────────────
 
 impl ACPIDeviceManagerState {
@@ -178,15 +117,9 @@ impl ACPIDeviceManagerState {
 
 // ───────────────────────────────────────────────────────────────────
 // VM state (aarch64, v1.14: adds resource_allocator)
+// VmState is the runtime type (re-exported above); conversion from v1.12 is here.
 // ───────────────────────────────────────────────────────────────────
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct VmState {
-    pub memory: GuestMemoryState,
-    pub gic: GicState,
-    pub resource_allocator: ResourceAllocator,
-}
-
 impl VmState {
     pub(crate) fn from(
         old_state: v1_12::VmState,

src/vmm/src/persist/v1_14/mod.rs

@@ -34,14 +34,18 @@ pub use aarch64::*;
 use crate::arch::{
     FIRST_ADDR_PAST_64BITS_MMIO, GSI_LEGACY_END, GSI_LEGACY_START, GSI_MSI_END, GSI_MSI_START,
     MEM_32BIT_DEVICES_SIZE, MEM_32BIT_DEVICES_START, MEM_64BIT_DEVICES_SIZE,
-    MEM_64BIT_DEVICES_START, PAST_64BITS_MMIO_SIZE, SYSTEM_MEM_SIZE, SYSTEM_MEM_START, VmState,
+    MEM_64BIT_DEVICES_START, PAST_64BITS_MMIO_SIZE, SYSTEM_MEM_SIZE, SYSTEM_MEM_START,
 };
+#[cfg(target_arch = "x86_64")]
+use crate::arch::VmState;
 use crate::device_manager::DevicesState;
 use crate::device_manager::mmio::MMIODeviceInfo;
 use crate::device_manager::pci_mngr::PciDevicesState;
 use crate::device_manager::persist::{
     ACPIDeviceManagerState, DeviceStates, MmdsState, VirtioDeviceState as ConnectedDeviceState,
 };
+#[cfg(target_arch = "aarch64")]
+use crate::device_manager::persist::ConnectedLegacyState;
 use crate::devices::acpi::vmgenid::VMGENID_MEM_SIZE;
 use crate::devices::virtio::balloon::device::HintingState;
 use crate::devices::virtio::balloon::persist::{BalloonState, BalloonStatsState};