diff --git a/exec-memory.h b/exec-memory.h
index 6707e40..ac1d07d 100644
--- a/exec-memory.h
+++ b/exec-memory.h
@@ -33,6 +33,9 @@ MemoryRegion *get_system_memory(void);
  */
 MemoryRegion *get_system_io(void);
 
+extern AddressSpace address_space_memory;
+extern AddressSpace address_space_io;
+
 #endif
 
 #endif
diff --git a/exec.c b/exec.c
index cfb8200..f561fba 100644
--- a/exec.c
+++ b/exec.c
@@ -116,8 +116,8 @@
 static MemoryRegion *system_memory;
 static MemoryRegion *system_io;
 
-static AddressSpace address_space_io;
-static AddressSpace address_space_memory;
+AddressSpace address_space_io;
+AddressSpace address_space_memory;
 
 MemoryRegion io_mem_ram, io_mem_rom, io_mem_unassigned, io_mem_notdirty;
 static MemoryRegion io_mem_subpage_ram;
@@ -3257,9 +3257,9 @@ static void memory_map_init(void)
     address_space_init(&address_space_io, system_io);
     address_space_io.name = "I/O";
 
-    memory_listener_register(&core_memory_listener, system_memory);
-    memory_listener_register(&io_memory_listener, system_io);
-    memory_listener_register(&tcg_memory_listener, system_memory);
+    memory_listener_register(&core_memory_listener, &address_space_memory);
+    memory_listener_register(&io_memory_listener, &address_space_io);
+    memory_listener_register(&tcg_memory_listener, &address_space_memory);
 }
 
 MemoryRegion *get_system_memory(void)
diff --git a/hw/vfio_pci.c b/hw/vfio_pci.c
index fb31f8e..e9399a1 100644
--- a/hw/vfio_pci.c
+++ b/hw/vfio_pci.c
@@ -1405,8 +1405,7 @@ static int vfio_connect_container(VFIOGroup *group)
         container->iommu_data.listener = vfio_memory_listener;
         container->iommu_data.release = vfio_listener_release;
 
-        memory_listener_register(&container->iommu_data.listener,
-                                 get_system_memory());
+        memory_listener_register(&container->iommu_data.listener, &address_space_memory);
     } else {
         error_report("vfio: No available IOMMU models\n");
         g_free(container);
diff --git a/hw/vhost.c b/hw/vhost.c
index 100f765..0b4ac3f 100644
--- a/hw/vhost.c
+++ b/hw/vhost.c
@@ -792,7 +792,7 @@ int vhost_dev_init(struct vhost_dev *hdev, int devfd, const char *devpath,
     hdev->log_size = 0;
     hdev->log_enabled = false;
     hdev->started = false;
-    memory_listener_register(&hdev->memory_listener, get_system_memory());
+    memory_listener_register(&hdev->memory_listener, &address_space_memory);
     hdev->force = force;
     return 0;
 fail:
diff --git a/hw/xen_pt.c b/hw/xen_pt.c
index 6f16182..8969360 100644
--- a/hw/xen_pt.c
+++ b/hw/xen_pt.c
@@ -746,8 +746,8 @@ static int xen_pt_initfn(PCIDevice *d)
     }
 
 out:
-    memory_listener_register(&s->memory_listener, get_system_memory());
-    memory_listener_register(&s->io_listener, get_system_io());
+    memory_listener_register(&s->memory_listener, &address_space_memory);
+    memory_listener_register(&s->io_listener, &address_space_io);
     XEN_PT_LOG(d, "Real physical device %02x:%02x.%d registered successfuly!\n",
                bus, slot, func);
 
diff --git a/kvm-all.c b/kvm-all.c
index 0a3c4ae..7b73df0 100644
--- a/kvm-all.c
+++ b/kvm-all.c
@@ -1375,8 +1375,8 @@ int kvm_init(void)
     }
 
     kvm_state = s;
-    memory_listener_register(&kvm_memory_listener, get_system_memory());
-    memory_listener_register(&kvm_io_listener, get_system_io());
+    memory_listener_register(&kvm_memory_listener, &address_space_memory);
+    memory_listener_register(&kvm_io_listener, &address_space_io);
 
     s->many_ioeventfds = kvm_check_many_ioeventfds();
 
diff --git a/memory.c b/memory.c
index 49a6ecc..0cf0177 100644
--- a/memory.c
+++ b/memory.c
@@ -147,7 +147,7 @@ static bool memory_listener_match(MemoryListener *listener,
 #define MEMORY_LISTENER_UPDATE_REGION(fr, as, dir, callback)            \
     MEMORY_LISTENER_CALL(callback, dir, (&(MemoryRegionSection) {       \
         .mr = (fr)->mr,                                                 \
-        .address_space = (as)->root,                                    \
+        .address_space = (as),                                          \
         .offset_within_region = (fr)->offset_in_region,                 \
         .size = int128_get64((fr)->addr.size),                          \
         .offset_within_address_space = int128_get64((fr)->addr.start),  \
@@ -593,7 +593,7 @@ static void address_space_add_del_ioeventfds(AddressSpace *as,
                                                   fds_new[inew]))) {
             fd = &fds_old[iold];
             section = (MemoryRegionSection) {
-                .address_space = as->root,
+                .address_space = as,
                 .offset_within_address_space = int128_get64(fd->addr.start),
                 .size = int128_get64(fd->addr.size),
             };
@@ -606,7 +606,7 @@ static void address_space_add_del_ioeventfds(AddressSpace *as,
                                                          fds_old[iold]))) {
             fd = &fds_new[inew];
             section = (MemoryRegionSection) {
-                .address_space = as->root,
+                .address_space = as,
                 .offset_within_address_space = int128_get64(fd->addr.start),
                 .size = int128_get64(fd->addr.size),
             };
@@ -1137,7 +1137,7 @@ static void memory_region_update_coalesced_range_as(MemoryRegion *mr, AddressSpa
     FOR_EACH_FLAT_RANGE(fr, as->current_map) {
         if (fr->mr == mr) {
             section = (MemoryRegionSection) {
-                .address_space = as->root,
+                .address_space = as,
                 .offset_within_address_space = int128_get64(fr->addr.start),
                 .size = int128_get64(fr->addr.size),
             };
@@ -1476,7 +1476,7 @@ static void listener_add_address_space(MemoryListener *listener,
     FlatRange *fr;
 
     if (listener->address_space_filter
-        && listener->address_space_filter != as->root) {
+        && listener->address_space_filter != as) {
         return;
     }
 
@@ -1489,7 +1489,7 @@ static void listener_add_address_space(MemoryListener *listener,
     FOR_EACH_FLAT_RANGE(fr, as->current_map) {
         MemoryRegionSection section = {
             .mr = fr->mr,
-            .address_space = as->root,
+            .address_space = as,
             .offset_within_region = fr->offset_in_region,
             .size = int128_get64(fr->addr.size),
             .offset_within_address_space = int128_get64(fr->addr.start),
@@ -1501,7 +1501,7 @@ static void listener_add_address_space(MemoryListener *listener,
     }
 }
 
-void memory_listener_register(MemoryListener *listener, MemoryRegion *filter)
+void memory_listener_register(MemoryListener *listener, AddressSpace *filter)
 {
     MemoryListener *other = NULL;
     AddressSpace *as;
diff --git a/memory.h b/memory.h
index 64d2b34..f5a13a4 100644
--- a/memory.h
+++ b/memory.h
@@ -187,7 +187,7 @@ typedef struct MemoryRegionSection MemoryRegionSection;
  */
 struct MemoryRegionSection {
     MemoryRegion *mr;
-    MemoryRegion *address_space;
+    AddressSpace *address_space;
     target_phys_addr_t offset_within_region;
     uint64_t size;
     target_phys_addr_t offset_within_address_space;
@@ -223,7 +223,7 @@ struct MemoryListener {
                                target_phys_addr_t addr, target_phys_addr_t len);
     /* Lower = earlier (during add), later (during del) */
     unsigned priority;
-    MemoryRegion *address_space_filter;
+    AddressSpace *address_space_filter;
     QTAILQ_ENTRY(MemoryListener) link;
 };
 
@@ -774,7 +774,7 @@ void memory_region_transaction_commit(void);
  * @listener: an object containing the callbacks to be called
  * @filter: if non-%NULL, only regions in this address space will be observed
  */
-void memory_listener_register(MemoryListener *listener, MemoryRegion *filter);
+void memory_listener_register(MemoryListener *listener, AddressSpace *filter);
 
 /**
  * memory_listener_unregister: undo the effect of memory_listener_register()
diff --git a/xen-all.c b/xen-all.c
index cf0e0c3..ea95c98 100644
--- a/xen-all.c
+++ b/xen-all.c
@@ -1116,7 +1116,7 @@ int xen_hvm_init(void)
 
     state->memory_listener = xen_memory_listener;
     QLIST_INIT(&state->physmap);
-    memory_listener_register(&state->memory_listener, get_system_memory());
+    memory_listener_register(&state->memory_listener, &address_space_memory);
     state->log_for_dirtybit = NULL;
 
     /* Initialize backend core & drivers */
