1 files changed, 19 insertions, 14 deletions

diff --git a/hw/virtio/virtio-balloon.c b/hw/virtio/virtio-balloon.c
index 2a2e58a297..2c30b3d8bd 100644
--- a/hw/virtio/virtio-balloon.c
+++ b/hw/virtio/virtio-balloon.c
@@ -31,6 +31,7 @@
 #endif
 
 #include "hw/virtio/virtio-bus.h"
+#include "hw/virtio/virtio-access.h"
 
 static void balloon_page(void *addr, int deflate)
 {
@@ -206,8 +207,9 @@ static void virtio_balloon_handle_output(VirtIODevice *vdev, VirtQueue *vq)
         while (iov_to_buf(elem.out_sg, elem.out_num, offset, &pfn, 4) == 4) {
             ram_addr_t pa;
             ram_addr_t addr;
+            int p = virtio_ldl_p(vdev, &pfn);
 
-            pa = (ram_addr_t)ldl_p(&pfn) << VIRTIO_BALLOON_PFN_SHIFT;
+            pa = (ram_addr_t) p << VIRTIO_BALLOON_PFN_SHIFT;
             offset += 4;
 
             /* FIXME: remove get_system_memory(), but how? */
@@ -248,8 +250,8 @@ static void virtio_balloon_receive_stats(VirtIODevice *vdev, VirtQueue *vq)
 
     while (iov_to_buf(elem->out_sg, elem->out_num, offset, &stat, sizeof(stat))
            == sizeof(stat)) {
-        uint16_t tag = tswap16(stat.tag);
-        uint64_t val = tswap64(stat.val);
+        uint16_t tag = virtio_tswap16(vdev, stat.tag);
+        uint64_t val = virtio_tswap64(vdev, stat.val);
 
         offset += sizeof(stat);
         if (tag < VIRTIO_BALLOON_S_NR)
@@ -325,10 +327,12 @@ static void virtio_balloon_to_target(void *opaque, ram_addr_t target)
 
 static void virtio_balloon_save(QEMUFile *f, void *opaque)
 {
-    VirtIOBalloon *s = VIRTIO_BALLOON(opaque);
-    VirtIODevice *vdev = VIRTIO_DEVICE(s);
+    virtio_save(VIRTIO_DEVICE(opaque), f);
+}
 
-    virtio_save(vdev, f);
+static void virtio_balloon_save_device(VirtIODevice *vdev, QEMUFile *f)
+{
+    VirtIOBalloon *s = VIRTIO_BALLOON(vdev);
 
     qemu_put_be32(f, s->num_pages);
     qemu_put_be32(f, s->actual);
@@ -336,17 +340,16 @@ static void virtio_balloon_save(QEMUFile *f, void *opaque)
 
 static int virtio_balloon_load(QEMUFile *f, void *opaque, int version_id)
 {
-    VirtIOBalloon *s = VIRTIO_BALLOON(opaque);
-    VirtIODevice *vdev = VIRTIO_DEVICE(s);
-    int ret;
-
     if (version_id != 1)
         return -EINVAL;
 
-    ret = virtio_load(vdev, f);
-    if (ret) {
-        return ret;
-    }
+    return virtio_load(VIRTIO_DEVICE(opaque), f, version_id);
+}
+
+static int virtio_balloon_load_device(VirtIODevice *vdev, QEMUFile *f,
+                                      int version_id)
+{
+    VirtIOBalloon *s = VIRTIO_BALLOON(vdev);
 
     s->num_pages = qemu_get_be32(f);
     s->actual = qemu_get_be32(f);
@@ -416,6 +419,8 @@ static void virtio_balloon_class_init(ObjectClass *klass, void *data)
     vdc->get_config = virtio_balloon_get_config;
     vdc->set_config = virtio_balloon_set_config;
     vdc->get_features = virtio_balloon_get_features;
+    vdc->save = virtio_balloon_save_device;
+    vdc->load = virtio_balloon_load_device;
 }
 
 static const TypeInfo virtio_balloon_info = {