Avoid repition in walkpgdir

2 files changed, 24 insertions, 60 deletions

diff --git a/mmu.h b/mmu.h
index 9450d90..df11806 100644
--- a/mmu.h
+++ b/mmu.h
@@ -91,26 +91,22 @@ struct segdesc {
 // +------+-------+--------------+----------+------------+-------------+
 //       \-PMX(va)-/\-PDPX(va)--/ \-PDX(va)-/ \-PTX(va)-/
 
-#define PMX(va)         (((uint64)(va) >> PML4XSHIFT) & PXMASK)
-#define PDPX(va)         (((uint64)(va) >> PDPXSHIFT) & PXMASK)
-// page directory index
-#define PDX(va)         (((uint64)(va) >> PDXSHIFT) & PXMASK)
-// page table index
-#define PTX(va)         (((uint64)(va) >> PTXSHIFT) & PXMASK)
-
-// construct virtual address from indexes and offset
-#define PGADDR(d, t, o) ((uint64)((d) << PDXSHIFT | (t) << PTXSHIFT | (o)))
-
 // Page directory and page table constants.
 #define NPDENTRIES      512    // # directory entries per page directory
 #define NPTENTRIES      512    // # PTEs per page table
 #define PGSIZE          4096    // bytes mapped by a page
-
-#define PTXSHIFT        12      // offset of PTX in a linear address
+#define PGSHIFT         12      // offset of PTX in a linear address
 #define PDXSHIFT        21      // offset of PDX in a linear address
-#define PDPXSHIFT       30      // offset of PDPX in a linear address
-#define PML4XSHIFT      39      // offset of PML4X in a linear address
-#define PXMASK          0X1FF
+
+#define PXMASK          0x1FF
+#define PXSHIFT(n)	(PGSHIFT+(9*(n)))
+#define PDX(va)         (((uint64)(va) >> PDXSHIFT) & PXMASK)
+#define PTX(va)         (((uint64)(va) >> PGSHIFT) & PXMASK)
+#define PX(n, va)	((((uint64) (va)) >> PXSHIFT(n)) & PXMASK)
+#define L_PML4           3
+    
+// construct virtual address from indexes and offset
+#define PGADDR(d, t, o) ((uint64)((d) << PDXSHIFT | (t) << PGSHIFT | (o)))
 
 #define PGROUNDUP(sz)  (((sz)+PGSIZE-1) & ~(PGSIZE-1))
 #define PGROUNDDOWN(a) (((a)) & ~(PGSIZE-1))
diff --git a/vm.c b/vm.c
index fb0cc33..4a0c446 100644
--- a/vm.c
+++ b/vm.c
@@ -66,53 +66,21 @@ seginit(void)
 static pte_t *
 walkpgdir(pde_t *pml4, const void *va, int alloc)
 {
-  pml4e_t *pml4e;
-  pdpe_t *pdp;
-  pdpe_t *pdpe;
+  pde_t *pgtab = pml4;
   pde_t *pde;
-  pde_t *pd;
-  pte_t *pgtab;
-
-  // level 4
-  pml4e = &pml4[PMX(va)];
-  if(*pml4e & PTE_P)
-    pdp = (pdpe_t*)P2V(PTE_ADDR(*pml4e));  
-  else {
-    if(!alloc || (pdp = (pdpe_t*)kalloc()) == 0)
-      return 0;
-    // Make sure all those PTE_P bits are zero.
-    memset(pdp, 0, PGSIZE);
-    // The permissions here are overly generous, but they can
-    // be further restricted by the permissions in the page table
-    // entries, if necessary.
-    *pml4e = V2P(pdp) | PTE_P | PTE_W | PTE_U;
-  }
-
-  // XXX avoid repetition
-
-  // level 3
-  pdpe = &pdp[PDPX(va)];  
-  if(*pdpe & PTE_P) 
-    pd = (pde_t*)P2V(PTE_ADDR(*pdpe));
-  else {
-    if(!alloc || (pd = (pde_t*)kalloc()) == 0)
-      return 0;
-    memset(pd, 0, PGSIZE);
-    *pdpe = V2P(pd) | PTE_P | PTE_W | PTE_U;
-  }
-
-  // level 2
-  pde = &pd[PDX(va)]; 
-  if(*pde & PTE_P)
-    pgtab = (pte_t*)P2V(PTE_ADDR(*pde));
-  else {
-    if(!alloc || (pgtab = (pte_t*)kalloc()) == 0)
-      return 0;
-    memset(pgtab, 0, PGSIZE);
-    *pde = V2P(pgtab) | PTE_P | PTE_W | PTE_U;
+  int level;
+  
+  for (level = L_PML4; level > 0; level--) {
+    pde = &pgtab[PX(level, va)];
+    if(*pde & PTE_P)
+      pgtab = (pte_t*)P2V(PTE_ADDR(*pde));
+    else {
+      if(!alloc || (pgtab = (pte_t*)kalloc()) == 0)
+        return 0;
+      memset(pgtab, 0, PGSIZE);
+      *pde = V2P(pgtab) | PTE_P | PTE_W | PTE_U;
+    }
   }
-
-  // level 1
   return &pgtab[PTX(va)];
 }