/*
  * Write-protect on the specified @sptep, @pt_protect indicates whether
- * spte write-protection is caused by protecting shadow page table.
+ * spte writ-protection is caused by protecting shadow page table.
+ * @flush indicates whether tlb need be flushed.
  *
  * Note: write protection is difference between drity logging and spte
  * protection:
  * - for spte protection, the spte can be writable only after unsync-ing
  *   shadow page.
  *
- * Return true if tlb need be flushed.
+ * Return true if the spte is dropped.
  */
-static bool spte_write_protect(struct kvm *kvm, u64 *sptep, bool pt_protect)
+static bool
+spte_write_protect(struct kvm *kvm, u64 *sptep, bool *flush, bool pt_protect)
 {
        u64 spte = *sptep;
 
 
        rmap_printk("rmap_write_protect: spte %p %llx\n", sptep, *sptep);
 
+       if (__drop_large_spte(kvm, sptep)) {
+               *flush |= true;
+               return true;
+       }
+
        if (pt_protect)
                spte &= ~SPTE_MMU_WRITEABLE;
        spte = spte & ~PT_WRITABLE_MASK;
 
-       return mmu_spte_update(sptep, spte);
+       *flush |= mmu_spte_update(sptep, spte);
+       return false;
 }
 
 static bool __rmap_write_protect(struct kvm *kvm, unsigned long *rmapp,
 
        for (sptep = rmap_get_first(*rmapp, &iter); sptep;) {
                BUG_ON(!(*sptep & PT_PRESENT_MASK));
+               if (spte_write_protect(kvm, sptep, &flush, pt_protect)) {
+                       sptep = rmap_get_first(*rmapp, &iter);
+                       continue;
+               }
 
-               flush |= spte_write_protect(kvm, sptep, pt_protect);
                sptep = rmap_get_next(&iter);
        }
 
                        break;
                }
 
-               drop_large_spte(vcpu, iterator.sptep);
-
                if (!is_shadow_present_pte(*iterator.sptep)) {
                        u64 base_addr = iterator.addr;