pt_element_t *table;
        pt_element_t *ptep;
        pt_element_t inherited_ar;
+       gfn_t gfn;
 };
 
 /*
                if (is_present_pte(*ptep) && !(*ptep &  PT_ACCESSED_MASK))
                        *ptep |= PT_ACCESSED_MASK;
 
-               if (!is_present_pte(*ptep) ||
-                   walker->level == PT_PAGE_TABLE_LEVEL ||
-                   (walker->level == PT_DIRECTORY_LEVEL &&
-                    (*ptep & PT_PAGE_SIZE_MASK) &&
-                    (PTTYPE == 64 || is_pse(vcpu))))
+               if (!is_present_pte(*ptep))
+                       break;
+
+               if (walker->level == PT_PAGE_TABLE_LEVEL) {
+                       walker->gfn = (*ptep & PT_BASE_ADDR_MASK)
+                               >> PAGE_SHIFT;
+                       break;
+               }
+
+               if (walker->level == PT_DIRECTORY_LEVEL
+                   && (*ptep & PT_PAGE_SIZE_MASK)
+                   && (PTTYPE == 64 || is_pse(vcpu))) {
+                       walker->gfn = (*ptep & PT_DIR_BASE_ADDR_MASK)
+                               >> PAGE_SHIFT;
+                       walker->gfn += PT_INDEX(addr, PT_PAGE_TABLE_LEVEL);
                        break;
+               }
 
                if (walker->level != 3 || is_long_mode(vcpu))
                        walker->inherited_ar &= walker->table[index];
 }
 
 static void FNAME(set_pte)(struct kvm_vcpu *vcpu, u64 guest_pte,
-                          u64 *shadow_pte, u64 access_bits)
+                          u64 *shadow_pte, u64 access_bits, gfn_t gfn)
 {
        ASSERT(*shadow_pte == 0);
        access_bits &= guest_pte;
        *shadow_pte = (guest_pte & PT_PTE_COPY_MASK);
        set_pte_common(vcpu, shadow_pte, guest_pte & PT_BASE_ADDR_MASK,
-                      guest_pte & PT_DIRTY_MASK, access_bits);
+                      guest_pte & PT_DIRTY_MASK, access_bits, gfn);
 }
 
 static void FNAME(set_pde)(struct kvm_vcpu *vcpu, u64 guest_pde,
-                          u64 *shadow_pte, u64 access_bits,
-                          int index)
+                          u64 *shadow_pte, u64 access_bits, gfn_t gfn)
 {
        gpa_t gaddr;
 
        ASSERT(*shadow_pte == 0);
        access_bits &= guest_pde;
-       gaddr = (guest_pde & PT_DIR_BASE_ADDR_MASK) + PAGE_SIZE * index;
+       gaddr = (gpa_t)gfn << PAGE_SHIFT;
        if (PTTYPE == 32 && is_cpuid_PSE36())
                gaddr |= (guest_pde & PT32_DIR_PSE36_MASK) <<
                        (32 - PT32_DIR_PSE36_SHIFT);
        *shadow_pte = guest_pde & PT_PTE_COPY_MASK;
        set_pte_common(vcpu, shadow_pte, gaddr,
-                      guest_pde & PT_DIRTY_MASK, access_bits);
+                      guest_pde & PT_DIRTY_MASK, access_bits, gfn);
 }
 
 /*
                                        *prev_shadow_ent |= PT_SHADOW_PS_MARK;
                                FNAME(set_pde)(vcpu, *guest_ent, shadow_ent,
                                               walker->inherited_ar,
-                                         PT_INDEX(addr, PT_PAGE_TABLE_LEVEL));
+                                              walker->gfn);
                        } else {
                                ASSERT(walker->level == PT_PAGE_TABLE_LEVEL);
-                               FNAME(set_pte)(vcpu, *guest_ent, shadow_ent, walker->inherited_ar);
+                               FNAME(set_pte)(vcpu, *guest_ent, shadow_ent,
+                                              walker->inherited_ar,
+                                              walker->gfn);
                        }
                        return shadow_ent;
                }
                return 0;
        }
 
-       gfn = (*guest_ent & PT64_BASE_ADDR_MASK) >> PAGE_SHIFT;
+       gfn = walker->gfn;
        if (kvm_mmu_lookup_page(vcpu, gfn)) {
                pgprintk("%s: found shadow page for %lx, marking ro\n",
                         __FUNCTION__, gfn);