[PATCH 07/11] rust: migration: add high-level migration wrappers

[Paolo Bonzini <pbonzini@redhat.com>]

Instead of dealing with pre/post callbacks, allow devices to
implement a snapshot/restore mechanism; this has two main
advantages:

- it can be easily implemented via procedural macros

- there can be generic implementations to deal with various
  kinds of interior-mutable containers, from BqlRefCell to Mutex,
  so that C code does not see Rust concepts such as Mutex<>.

Using it is easy; you can implement the snapshot/restore trait
ToMigrationState and declare your state like:

     regs: Migratable<Mutex<MyDeviceRegisters>>

Migratable<> allows dereferencing to the underlying object with
no run-time cost.

Note that Migratable<> actually does not accept ToMigrationState,
only the similar ToMigrationStateShared trait that the user will mostly
not care about.  This is required by the fact that pre/post callbacks
take a &self, and ensures that the argument is a Mutex or BqlRefCell
(including an array or Arc<> thereof).

Signed-off-by: Paolo Bonzini <pbonzini@redhat.com>
---
 docs/devel/rust.rst              |   1 +
 rust/migration/meson.build       |   1 +
 rust/migration/src/lib.rs        |   3 +
 rust/migration/src/migratable.rs | 434 +++++++++++++++++++++++++++++++
 4 files changed, 439 insertions(+)
 create mode 100644 rust/migration/src/migratable.rs

diff --git a/docs/devel/rust.rst b/docs/devel/rust.rst
index 2f0ab2e2821..79c26d9d165 100644
--- a/docs/devel/rust.rst
+++ b/docs/devel/rust.rst
@@ -155,6 +155,7 @@ module                     status
 ``hwcore::irq``            complete
 ``hwcore::qdev``           stable
 ``hwcore::sysbus``         stable
+``migration::migratable``  proof of concept
 ``migration::vmstate``     stable
 ``qom``                    stable
 ``system::memory``         stable
diff --git a/rust/migration/meson.build b/rust/migration/meson.build
index e381c76d3e8..1ba0ec21d72 100644
--- a/rust/migration/meson.build
+++ b/rust/migration/meson.build
@@ -31,6 +31,7 @@ _migration_rs = static_library(
     [
       'src/lib.rs',
       'src/bindings.rs',
+      'src/migratable.rs',
       'src/vmstate.rs',
     ],
     {'.' : _migration_bindings_inc_rs},
diff --git a/rust/migration/src/lib.rs b/rust/migration/src/lib.rs
index 5f51dde4406..efe9896b619 100644
--- a/rust/migration/src/lib.rs
+++ b/rust/migration/src/lib.rs
@@ -2,5 +2,8 @@
 
 pub mod bindings;
 
+pub mod migratable;
+pub use migratable::*;
+
 pub mod vmstate;
 pub use vmstate::*;
diff --git a/rust/migration/src/migratable.rs b/rust/migration/src/migratable.rs
new file mode 100644
index 00000000000..46e533c16d1
--- /dev/null
+++ b/rust/migration/src/migratable.rs
@@ -0,0 +1,434 @@
+// Copyright 2025 Red Hat, Inc.
+// Author(s): Paolo Bonzini <pbonzini@redhat.com>
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+use std::{
+    fmt,
+    mem::size_of,
+    ptr::{self, addr_of, NonNull},
+    sync::{Arc, Mutex},
+};
+
+use bql::{BqlCell, BqlRefCell};
+use common::Zeroable;
+
+use crate::{
+    bindings, vmstate_fields_ref, vmstate_of, InvalidError, VMState, VMStateDescriptionBuilder,
+};
+
+/// Enables QEMU migration support even when a type is wrapped with
+/// synchronization primitives (like `Mutex`) that the C migration
+/// code cannot directly handle. The trait provides methods to
+/// extract essential state for migration and restore it after
+/// migration completes.
+///
+/// On top of extracting data from synchronization wrappers during save
+/// and restoring it during load, it's also possible to use `ToMigrationState`
+/// to convert runtime representations to migration-safe formats.
+///
+/// # Examples
+///
+/// ```
+/// use bql::BqlCell;
+/// use migration::{InvalidError, ToMigrationState, VMState};
+/// # use migration::VMStateField;
+///
+/// # #[derive(Debug, PartialEq, Eq)]
+/// struct DeviceState {
+///     counter: BqlCell<u32>,
+///     enabled: bool,
+/// }
+///
+/// # #[derive(Debug)]
+/// #[derive(Default)]
+/// struct DeviceMigrationState {
+///     counter: u32,
+///     enabled: bool,
+/// }
+///
+/// # unsafe impl VMState for DeviceMigrationState {
+/// #     const BASE: VMStateField = ::common::Zeroable::ZERO;
+/// # }
+/// impl ToMigrationState for DeviceState {
+///     type Migrated = DeviceMigrationState;
+///
+///     fn snapshot_migration_state(
+///         &self,
+///         target: &mut Self::Migrated,
+///     ) -> Result<(), InvalidError> {
+///         target.counter = self.counter.get();
+///         target.enabled = self.enabled;
+///         Ok(())
+///     }
+///
+///     fn restore_migrated_state_mut(
+///         &mut self,
+///         source: Self::Migrated,
+///         _version_id: u8,
+///     ) -> Result<(), InvalidError> {
+///         self.counter.set(source.counter);
+///         self.enabled = source.enabled;
+///         Ok(())
+///     }
+/// }
+/// # bql::start_test();
+/// # let dev = DeviceState { counter: 10.into(), enabled: true };
+/// # let mig = dev.to_migration_state().unwrap();
+/// # assert!(matches!(*mig, DeviceMigrationState { counter: 10, enabled: true }));
+/// # let mut dev2 = DeviceState { counter: 42.into(), enabled: false };
+/// # dev2.restore_migrated_state_mut(*mig, 1).unwrap();
+/// # assert_eq!(dev2, dev);
+/// ```
+pub trait ToMigrationState {
+    /// The type used to represent the migrated state.
+    type Migrated: Default + VMState;
+
+    /// Capture the current state into a migration-safe format, failing
+    /// if the state cannot be migrated.
+    fn snapshot_migration_state(&self, target: &mut Self::Migrated) -> Result<(), InvalidError>;
+
+    /// Restores state from a migrated representation, failing if the
+    /// state cannot be restored.
+    fn restore_migrated_state_mut(
+        &mut self,
+        source: Self::Migrated,
+        version_id: u8,
+    ) -> Result<(), InvalidError>;
+
+    /// Convenience method to combine allocation and state capture
+    /// into a single operation.
+    fn to_migration_state(&self) -> Result<Box<Self::Migrated>, InvalidError> {
+        let mut migrated = Box::<Self::Migrated>::default();
+        self.snapshot_migration_state(&mut migrated)?;
+        Ok(migrated)
+    }
+}
+
+// Implementations for primitive types.  Do not use a blanket implementation
+// for all Copy types, because [T; N] is Copy if T is Copy; that would conflict
+// with the below implementation for arrays.
+macro_rules! impl_for_primitive {
+    ($($t:ty),*) => {
+        $(
+            impl ToMigrationState for $t {
+                type Migrated = Self;
+
+                fn snapshot_migration_state(
+                    &self,
+                    target: &mut Self::Migrated,
+                ) -> Result<(), InvalidError> {
+                    *target = *self;
+                    Ok(())
+                }
+
+                fn restore_migrated_state_mut(
+                    &mut self,
+                    source: Self::Migrated,
+                    _version_id: u8,
+                ) -> Result<(), InvalidError> {
+                    *self = source;
+                    Ok(())
+                }
+            }
+        )*
+    };
+}
+
+impl_for_primitive!(u8, u16, u32, u64, i8, i16, i32, i64, bool);
+
+impl<T: ToMigrationState, const N: usize> ToMigrationState for [T; N]
+where
+    [T::Migrated; N]: Default,
+{
+    type Migrated = [T::Migrated; N];
+
+    fn snapshot_migration_state(&self, target: &mut Self::Migrated) -> Result<(), InvalidError> {
+        for (item, target_item) in self.iter().zip(target.iter_mut()) {
+            item.snapshot_migration_state(target_item)?;
+        }
+        Ok(())
+    }
+
+    fn restore_migrated_state_mut(
+        &mut self,
+        source: Self::Migrated,
+        version_id: u8,
+    ) -> Result<(), InvalidError> {
+        for (item, source_item) in self.iter_mut().zip(source) {
+            item.restore_migrated_state_mut(source_item, version_id)?;
+        }
+        Ok(())
+    }
+}
+
+impl<T: ToMigrationState> ToMigrationState for Mutex<T> {
+    type Migrated = T::Migrated;
+
+    fn snapshot_migration_state(&self, target: &mut Self::Migrated) -> Result<(), InvalidError> {
+        self.lock().unwrap().snapshot_migration_state(target)
+    }
+
+    fn restore_migrated_state_mut(
+        &mut self,
+        source: Self::Migrated,
+        version_id: u8,
+    ) -> Result<(), InvalidError> {
+        self.get_mut()
+            .unwrap()
+            .restore_migrated_state_mut(source, version_id)
+    }
+}
+
+impl<T: ToMigrationState> ToMigrationState for BqlRefCell<T> {
+    type Migrated = T::Migrated;
+
+    fn snapshot_migration_state(&self, target: &mut Self::Migrated) -> Result<(), InvalidError> {
+        self.borrow().snapshot_migration_state(target)
+    }
+
+    fn restore_migrated_state_mut(
+        &mut self,
+        source: Self::Migrated,
+        version_id: u8,
+    ) -> Result<(), InvalidError> {
+        self.get_mut()
+            .restore_migrated_state_mut(source, version_id)
+    }
+}
+
+/// Extension trait for types that support migration state restoration
+/// through interior mutability.
+///
+/// This trait extends [`ToMigrationState`] for types that can restore
+/// their state without requiring mutable access. While user structs
+/// will generally use `ToMigrationState`, the device will have multiple
+/// references and therefore the device struct has to employ an interior
+/// mutability wrapper like [`Mutex`] or [`BqlRefCell`].
+///
+/// Anything that implements this trait can in turn be used within
+/// [`Migratable<T>`], which makes no assumptions on how to achieve mutable
+/// access to the runtime state.
+///
+/// # Examples
+///
+/// ```
+/// use std::sync::Mutex;
+///
+/// use migration::ToMigrationStateShared;
+///
+/// let device_state = Mutex::new(42);
+/// // Can restore without &mut access
+/// device_state.restore_migrated_state(100, 1).unwrap();
+/// assert_eq!(*device_state.lock().unwrap(), 100);
+/// ```
+pub trait ToMigrationStateShared: ToMigrationState {
+    /// Restores state from a migrated representation to an interior-mutable
+    /// object.  Similar to `restore_migrated_state_mut`, but requires a
+    /// shared reference; therefore it can be used to restore a device's
+    /// state even though devices have multiple references to them.
+    fn restore_migrated_state(
+        &self,
+        source: Self::Migrated,
+        version_id: u8,
+    ) -> Result<(), InvalidError>;
+}
+
+impl<T: ToMigrationStateShared, const N: usize> ToMigrationStateShared for [T; N]
+where
+    [T::Migrated; N]: Default,
+{
+    fn restore_migrated_state(
+        &self,
+        source: Self::Migrated,
+        version_id: u8,
+    ) -> Result<(), InvalidError> {
+        for (item, source_item) in self.iter().zip(source) {
+            item.restore_migrated_state(source_item, version_id)?;
+        }
+        Ok(())
+    }
+}
+
+// Arc requires the contained object to be interior-mutable
+impl<T: ToMigrationStateShared> ToMigrationState for Arc<T> {
+    type Migrated = T::Migrated;
+
+    fn snapshot_migration_state(&self, target: &mut Self::Migrated) -> Result<(), InvalidError> {
+        (**self).snapshot_migration_state(target)
+    }
+
+    fn restore_migrated_state_mut(
+        &mut self,
+        source: Self::Migrated,
+        version_id: u8,
+    ) -> Result<(), InvalidError> {
+        (**self).restore_migrated_state(source, version_id)
+    }
+}
+
+impl<T: ToMigrationStateShared> ToMigrationStateShared for Arc<T> {
+    fn restore_migrated_state(
+        &self,
+        source: Self::Migrated,
+        version_id: u8,
+    ) -> Result<(), InvalidError> {
+        (**self).restore_migrated_state(source, version_id)
+    }
+}
+
+// Interior-mutable types.  Note how they only require ToMigrationState for
+// the inner type!
+
+impl<T: ToMigrationState> ToMigrationStateShared for Mutex<T> {
+    fn restore_migrated_state(
+        &self,
+        source: Self::Migrated,
+        version_id: u8,
+    ) -> Result<(), InvalidError> {
+        self.lock()
+            .unwrap()
+            .restore_migrated_state_mut(source, version_id)
+    }
+}
+
+impl<T: ToMigrationState> ToMigrationStateShared for BqlRefCell<T> {
+    fn restore_migrated_state(
+        &self,
+        source: Self::Migrated,
+        version_id: u8,
+    ) -> Result<(), InvalidError> {
+        self.borrow_mut()
+            .restore_migrated_state_mut(source, version_id)
+    }
+}
+
+/// A wrapper that enables QEMU migration for types with shared state.
+///
+/// `Migratable<T>` provides a bridge between Rust types that use interior
+/// mutability (like `Mutex<T>`) and QEMU's C-based migration infrastructure.
+/// It manages the lifecycle of migration state and provides automatic
+/// conversion between runtime and migration representations.
+///
+/// ```ignore
+/// # use std::sync::Mutex;
+/// # use migration::Migratable;
+///
+/// pub struct DeviceRegs {
+///     status: u32,
+/// }
+///
+/// pub struct SomeDevice {
+///     // ...
+///     registers: Migratable<Mutex<DeviceRegs>>,
+/// }
+/// ```
+#[repr(C)]
+pub struct Migratable<T: ToMigrationStateShared> {
+    /// Pointer to migration state, valid only during migration operations.
+    /// C vmstate does not support NULL pointers, so no `Option<Box<>>`.
+    migration_state: BqlCell<*mut T::Migrated>,
+
+    /// The runtime state that can be accessed during normal operation
+    runtime_state: T,
+}
+
+impl<T: ToMigrationStateShared> std::ops::Deref for Migratable<T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        &self.runtime_state
+    }
+}
+
+impl<T: ToMigrationStateShared> std::ops::DerefMut for Migratable<T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.runtime_state
+    }
+}
+
+impl<T: ToMigrationStateShared> Migratable<T> {
+    /// Creates a new `Migratable` wrapper around the given runtime state.
+    ///
+    /// # Returns
+    /// A new `Migratable` instance ready for use and migration
+    pub fn new(runtime_state: T) -> Self {
+        Self {
+            migration_state: BqlCell::new(ptr::null_mut()),
+            runtime_state,
+        }
+    }
+
+    fn pre_save(&self) -> Result<(), InvalidError> {
+        let state = self.runtime_state.to_migration_state()?;
+        self.migration_state.set(Box::into_raw(state));
+        Ok(())
+    }
+
+    fn post_save(&self) -> Result<(), InvalidError> {
+        let state = unsafe { Box::from_raw(self.migration_state.replace(ptr::null_mut())) };
+        drop(state);
+        Ok(())
+    }
+
+    fn pre_load(&self) -> Result<(), InvalidError> {
+        self.migration_state
+            .set(Box::into_raw(Box::<T::Migrated>::default()));
+        Ok(())
+    }
+
+    fn post_load(&self, version_id: u8) -> Result<(), InvalidError> {
+        let state = unsafe { Box::from_raw(self.migration_state.replace(ptr::null_mut())) };
+        self.runtime_state
+            .restore_migrated_state(*state, version_id)
+    }
+}
+
+impl<T: ToMigrationStateShared + fmt::Debug> fmt::Debug for Migratable<T>
+where
+    T::Migrated: fmt::Debug,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let mut struct_f = f.debug_struct("Migratable");
+        struct_f.field("runtime_state", &self.runtime_state);
+
+        let state = NonNull::new(self.migration_state.get()).map(|x| unsafe { x.as_ref() });
+        struct_f.field("migration_state", &state);
+        struct_f.finish()
+    }
+}
+
+impl<T: ToMigrationStateShared + Default> Default for Migratable<T> {
+    fn default() -> Self {
+        Self::new(T::default())
+    }
+}
+
+impl<T: 'static + ToMigrationStateShared> Migratable<T> {
+    const FIELD: bindings::VMStateField = vmstate_of!(Self, migration_state);
+
+    const FIELDS: &[bindings::VMStateField] = vmstate_fields_ref! {
+        Migratable::<T>::FIELD
+    };
+
+    const VMSD: &'static bindings::VMStateDescription = VMStateDescriptionBuilder::<Self>::new()
+        .version_id(1)
+        .minimum_version_id(1)
+        .pre_save(&Self::pre_save)
+        .pre_load(&Self::pre_load)
+        .post_save(&Self::post_save)
+        .post_load(&Self::post_load)
+        .fields(Self::FIELDS)
+        .build()
+        .as_ref();
+}
+
+unsafe impl<T: 'static + ToMigrationStateShared> VMState for Migratable<T> {
+    const BASE: bindings::VMStateField = {
+        bindings::VMStateField {
+            vmsd: addr_of!(*Self::VMSD),
+            size: size_of::<Self>(),
+            flags: bindings::VMStateFlags::VMS_STRUCT,
+            ..Zeroable::ZERO
+        }
+    };
+}
-- 
2.51.0