2 files changed, 29 insertions, 13 deletions

diff --git a/include/linux/kvm_host.h b/include/linux/kvm_host.h
index 1a371447fd45..193bca68372d 100644
--- a/include/linux/kvm_host.h
+++ b/include/linux/kvm_host.h
@@ -354,6 +354,7 @@ struct kvm_memslots {
         /* The mapping table from slot id to the index in memslots[]. */
         short id_to_index[KVM_MEM_SLOTS_NUM];
         atomic_t lru_slot;
+        int used_slots;
 };
 
 struct kvm {
@@ -791,19 +792,28 @@ static inline void kvm_guest_exit(void)
 static inline struct kvm_memory_slot *
 search_memslots(struct kvm_memslots *slots, gfn_t gfn)
 {
+        int start = 0, end = slots->used_slots;
         int slot = atomic_read(&slots->lru_slot);
-        struct kvm_memory_slot *memslot = &slots->memslots[slot];
-
-        if (gfn >= memslot->base_gfn &&
-            gfn < memslot->base_gfn + memslot->npages)
-                return memslot;
-
-        kvm_for_each_memslot(memslot, slots)
-                if (gfn >= memslot->base_gfn &&
-                      gfn < memslot->base_gfn + memslot->npages) {
-                        atomic_set(&slots->lru_slot, memslot - slots->memslots);
-                        return memslot;
-                }
+        struct kvm_memory_slot *memslots = slots->memslots;
+
+        if (gfn >= memslots[slot].base_gfn &&
+            gfn < memslots[slot].base_gfn + memslots[slot].npages)
+                return &memslots[slot];
+
+        while (start < end) {
+                slot = start + (end - start) / 2;
+
+                if (gfn >= memslots[slot].base_gfn)
+                        end = slot;
+                else
+                        start = slot + 1;
+        }
+
+        if (gfn >= memslots[start].base_gfn &&
+            gfn < memslots[start].base_gfn + memslots[start].npages) {
+                atomic_set(&slots->lru_slot, start);
+                return &memslots[start];
+        }
 
         return NULL;
 }
diff --git a/virt/kvm/kvm_main.c b/virt/kvm/kvm_main.c
index 162817f853ec..759af6596a07 100644
--- a/virt/kvm/kvm_main.c
+++ b/virt/kvm/kvm_main.c
@@ -679,8 +679,14 @@ static void update_memslots(struct kvm_memslots *slots,
         struct kvm_memory_slot *mslots = slots->memslots;
 
         WARN_ON(mslots[i].id != id);
-        if (!new->npages)
+        if (!new->npages) {
                 new->base_gfn = 0;
+                if (mslots[i].npages)
+                        slots->used_slots--;
+        } else {
+                if (!mslots[i].npages)
+                        slots->used_slots++;
+        }
 
         while (i < KVM_MEM_SLOTS_NUM - 1 &&
                new->base_gfn <= mslots[i + 1].base_gfn) {