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>
---
 rust/Cargo.lock                  |   1 +
 rust/migration/Cargo.toml        |   1 +
 rust/migration/meson.build       |   1 +
 rust/migration/src/lib.rs        |   3 +
 rust/migration/src/migratable.rs | 430 +++++++++++++++++++++++++++++++
 5 files changed, 436 insertions(+)
 create mode 100644 rust/migration/src/migratable.rs

diff --git a/rust/Cargo.lock b/rust/Cargo.lock
index 960f603cedb..32db90066f1 100644
--- a/rust/Cargo.lock
+++ b/rust/Cargo.lock
@@ -140,6 +140,7 @@ version = "0.1.0"
 dependencies = [
  "bql",
  "common",
+ "qemu_macros",
  "util",
 ]
 
diff --git a/rust/migration/Cargo.toml b/rust/migration/Cargo.toml
index f4a86275152..8efce7a72cb 100644
--- a/rust/migration/Cargo.toml
+++ b/rust/migration/Cargo.toml
@@ -15,6 +15,7 @@ rust-version.workspace = true
 [dependencies]
 bql = { path = "../bql" }
 common = { path = "../common" }
+qemu_macros = { path = "../qemu-macros" }
 util = { path = "../util" }
 
 [lints]
diff --git a/rust/migration/meson.build b/rust/migration/meson.build
index 2f38da9220f..c258881790d 100644
--- a/rust/migration/meson.build
+++ b/rust/migration/meson.build
@@ -30,6 +30,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..d09eeb35f11
--- /dev/null
+++ b/rust/migration/src/migratable.rs
@@ -0,0 +1,430 @@
+// 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(())
+    }
+}
+
+/// 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(())
+    }
+}
+
+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 only require ToMigrationState for the inner type!
+
+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.restore_migrated_state(source, version_id)
+    }
+}
+
+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> 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)
+    }
+}
+
+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.take()) };
+        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.take()) };
+        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.as_ptr())
+        .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

On Sat, Sep 20, 2025 at 04:29:55PM +0200, Paolo Bonzini wrote:
> Date: Sat, 20 Sep 2025 16:29:55 +0200
> From: Paolo Bonzini <pbonzini@redhat.com>
> Subject: [PATCH 4/7] rust: migration: add high-level migration wrappers
> X-Mailer: git-send-email 2.51.0
> 
> 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>
> ---
>  rust/Cargo.lock                  |   1 +
>  rust/migration/Cargo.toml        |   1 +
>  rust/migration/meson.build       |   1 +
>  rust/migration/src/lib.rs        |   3 +
>  rust/migration/src/migratable.rs | 430 +++++++++++++++++++++++++++++++
>  5 files changed, 436 insertions(+)
>  create mode 100644 rust/migration/src/migratable.rs

The entire wrapper is quite nice. So,

Reviewed-by: Zhao Liu <zhao1.liu@intel.com>


Only a few comments for comments inline:

> +// Interior-mutable types only require ToMigrationState for the inner type!
> +

extra blank line.

> +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)
> +    }

Or maybe your previous sentence is worth commenting on here:

// For non-BQL-protected device we cannot know that another
// thread isn't taking the lock. So, always acquire the lock.

> +    fn restore_migrated_state_mut(
> +        &mut self,
> +        source: Self::Migrated,
> +        version_id: u8,
> +    ) -> Result<(), InvalidError> {
> +        self.restore_migrated_state(source, version_id)
> +    }
> +}
> +
> +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)
> +    }
> +}
> +

On Thu, Sep 25, 2025, 10:43 Zhao Liu <zhao1.liu@intel.com> wrote:

> > +    fn snapshot_migration_state(&self, target: &mut Self::Migrated) ->
> Result<(), InvalidError> {
> > +        self.lock().unwrap().snapshot_migration_state(target)
> > +    }
>
> Or maybe your previous sentence is worth commenting on here:
>
> // For non-BQL-protected device we cannot know that another
> // thread isn't taking the lock. So, always acquire the lock.
>

I don't think here there's any alternative, that is a way to write code
without taking the lock. However...

>
> > +    fn restore_migrated_state_mut(
> > +        &mut self,
> > +        source: Self::Migrated,
> > +        version_id: u8,
> > +    ) -> Result<(), InvalidError> {
> > +        self.restore_migrated_state(source, version_id)
>

... this could use get_mut().

Paolo

> +    }
> > +}
> > +
> > +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)
> > +    }
> > +}
> > +
>
>