[4/6] ide: fix migration in the middle of pio operation

diff --git a/hw/ide/core.c b/hw/ide/core.c
index ccb7d32..db00083 100644
--- a/hw/ide/core.c
+++ b/hw/ide/core.c
@@ -2664,6 +2664,7 @@  static void ide_init1(IDEBus *bus, int unit)
     s->unit = unit;
     s->drive_serial = drive_serial++;
     s->io_buffer = qemu_blockalign(s->bs, IDE_DMA_BUF_SECTORS*512 + 4);
+    s->io_buffer_total_len = IDE_DMA_BUF_SECTORS*512 + 4;
     s->smart_selftest_data = qemu_blockalign(s->bs, 512);
     s->sector_write_timer = qemu_new_timer(vm_clock,
                                            ide_sector_write_timer_cb, s);
@@ -2726,6 +2727,26 @@  static bool is_identify_set(void *opaque, int version_id)
     return s->identify_set != 0;
 }

+static EndTransferFunc* transfer_end_table[] = {
+        ide_sector_read,
+        ide_sector_write,
+        ide_transfer_stop,
+        ide_atapi_cmd_reply_end,
+        ide_atapi_cmd,
+        ide_dummy_transfer_stop,
+};
+
+static int transfer_end_table_idx(EndTransferFunc *fn)
+{
+    int i;
+
+    for (i = 0; i < ARRAY_SIZE(transfer_end_table); i++)
+        if (transfer_end_table[i] == fn)
+            return i;
+
+    return -1;
+}
+
 static int ide_drive_post_load(void *opaque, int version_id)
 {
     IDEState *s = opaque;
@@ -2739,6 +2760,66 @@  static int ide_drive_post_load(void *opaque, int version_id)
     return 0;
 }

+static int ide_drive_pio_post_load(void *opaque, int version_id)
+{
+    IDEState *s = opaque;
+
+    if (s->end_transfer_fn_idx < 0 ||
+        s->end_transfer_fn_idx > ARRAY_SIZE(transfer_end_table)) {
+        return -EINVAL;
+    }
+    s->end_transfer_func = transfer_end_table[s->end_transfer_fn_idx];
+    s->data_ptr = s->io_buffer + s->cur_io_buffer_offset;
+    s->data_end = s->data_ptr + s->cur_io_buffer_len;
+
+    return 0;
+}
+
+static void ide_drive_pio_pre_save(void *opaque)
+{
+    IDEState *s = opaque;
+    int idx;
+
+    s->cur_io_buffer_offset = s->data_ptr - s->io_buffer;
+    s->cur_io_buffer_len = s->data_end - s->data_ptr;
+
+    idx = transfer_end_table_idx(s->end_transfer_func);
+    if (idx == -1) {
+        fprintf(stderr, "%s: invalid end_transfer_func for DRQ_STAT\n",
+                        __func__);
+        s->end_transfer_fn_idx = 2;
+    } else {
+        s->end_transfer_fn_idx = idx;
+    }
+}
+
+static bool ide_drive_pio_state_needed(void *opaque)
+{
+    IDEState *s = opaque;
+
+    return (s->status & DRQ_STAT) != 0;
+}
+
+const VMStateDescription vmstate_ide_drive_pio_state = {
+    .name = "ide_drive/pio_state",
+    .version_id = 1,
+    .minimum_version_id = 1,
+    .minimum_version_id_old = 1,
+    .pre_save = ide_drive_pio_pre_save,
+    .post_load = ide_drive_pio_post_load,
+    .fields      = (VMStateField []) {
+        VMSTATE_INT32(req_nb_sectors, IDEState),
+        VMSTATE_VARRAY_INT32(io_buffer, IDEState, io_buffer_total_len, 1,
+			     vmstate_info_uint8, uint8_t),
+        VMSTATE_INT32(cur_io_buffer_offset, IDEState),
+        VMSTATE_INT32(cur_io_buffer_len, IDEState),
+        VMSTATE_UINT8(end_transfer_fn_idx, IDEState),
+        VMSTATE_INT32(elementary_transfer_size, IDEState),
+        VMSTATE_INT32(packet_transfer_size, IDEState),
+        VMSTATE_END_OF_LIST()
+    }
+};
+
 const VMStateDescription vmstate_ide_drive = {
     .name = "ide_drive",
     .version_id = 3,
@@ -2766,8 +2847,15 @@  const VMStateDescription vmstate_ide_drive = {
         VMSTATE_UINT8(sense_key, IDEState),
         VMSTATE_UINT8(asc, IDEState),
         VMSTATE_UINT8_V(cdrom_changed, IDEState, 3),
-        /* XXX: if a transfer is pending, we do not save it yet */
         VMSTATE_END_OF_LIST()
+    },
+    .subsections = (VMStateSubsection []) {
+        {
+            .vmsd = &vmstate_ide_drive_pio_state,
+            .needed = ide_drive_pio_state_needed,
+        }, {
+            /* empty */
+        }
     }
 };

diff --git a/hw/ide/internal.h b/hw/ide/internal.h
index 75745ee..4165543 100644
--- a/hw/ide/internal.h
+++ b/hw/ide/internal.h
@@ -419,6 +419,11 @@  struct IDEState {
     uint8_t *data_ptr;
     uint8_t *data_end;
     uint8_t *io_buffer;
+    /* PIO save/restore */
+    int32_t io_buffer_total_len;
+    int cur_io_buffer_offset;
+    int cur_io_buffer_len;
+    uint8_t end_transfer_fn_idx;
     QEMUTimer *sector_write_timer; /* only used for win2k install hack */
     uint32_t irq_count; /* counts IRQs when using win2k install hack */
     /* CF-ATA extended error */