Improve pmap protect

include/aarch64/pmap.h

@@ -55,15 +55,44 @@ static __no_profile inline bool pde_valid_p(pde_t *pdep) {
 
 void *phys_to_dmap(paddr_t addr);
 
-static __no_profile inline pde_t *pde_ptr(paddr_t pd_pa, int lvl, vaddr_t va) {
-  pde_t *pde = phys_to_dmap(pd_pa);
-  if (lvl == 0)
-    return pde + L0_INDEX(va);
-  if (lvl == 1)
-    return pde + L1_INDEX(va);
-  if (lvl == 2)
-    return pde + L2_INDEX(va);
-  return pde + L3_INDEX(va);
+static __no_profile inline size_t pde_index(int lvl, vaddr_t va) {
+  switch (lvl) {
+    case 0:
+      return L0_INDEX(va);
+    case 1:
+      return L1_INDEX(va);
+    case 2:
+      return L2_INDEX(va);
+    case 3:
+      return L3_INDEX(va);
+  }
+  panic("Invalid level: %d", lvl);
+}
+
+static __no_profile inline bool pde_valid_index(int lvl, size_t index) {
+  switch (lvl) {
+    case 0:
+      return index < L0_ENTRIES;
+    case 1:
+    case 2:
+    case 3:
+      return index < Ln_ENTRIES;
+  }
+  panic("Invalid level: %d", lvl);
+}
+
+static __no_profile inline size_t pde_size(int lvl) {
+  switch (lvl) {
+    case 0:
+      return L0_SIZE;
+    case 1:
+      return L1_SIZE;
+    case 2:
+      return L2_SIZE;
+    case 3:
+      return L3_SIZE;
+  }
+  panic("Invalid level: %d", lvl);
 }
 
 /*

include/mips/pmap.h

@@ -83,11 +83,36 @@ static __no_profile inline bool pde_valid_p(pde_t *pdep) {
 
 void *phys_to_dmap(paddr_t addr);
 
-static __no_profile inline pde_t *pde_ptr(paddr_t pd_pa, int lvl, vaddr_t va) {
-  pde_t *pde = phys_to_dmap(pd_pa);
-  if (lvl == 0)
-    return pde + PDE_INDEX(va);
-  return pde + PTE_INDEX(va);
+static __no_profile inline size_t pde_index(int lvl, vaddr_t va) {
+  switch (lvl) {
+    case 0:
+      return PDE_INDEX(va);
+    case 1:
+      return PTE_INDEX(va);
+  }
+  panic("Invalid level: %d", lvl);
+}
+
+static __no_profile inline bool pde_valid_index(int lvl, size_t index) {
+  switch (lvl) {
+    case 0:
+      return index < PD_ENTRIES;
+    case 1:
+    case 2:
+    case 3:
+      return index < PT_ENTRIES;
+  }
+  panic("Invalid level: %d", lvl);
+}
+
+static __no_profile inline size_t pde_size(int lvl) {
+  switch (lvl) {
+    case 0:
+      return PAGESIZE * PT_ENTRIES;
+    case 1:
+      return PAGESIZE;
+  }
+  panic("Invalid level: %d", lvl);
 }
 
 /*

include/riscv/pmap.h

@@ -52,17 +52,44 @@ static __no_profile inline bool pde_valid_p(pde_t *pdep) {
 
 void *phys_to_dmap(paddr_t addr);
 
-static __no_profile inline pde_t *pde_ptr(paddr_t pd_pa, int lvl, vaddr_t va) {
-  pde_t *pde = phys_to_dmap(pd_pa);
-  if (lvl == 0)
-    return pde + L0_INDEX(va);
-#if __riscv_xlen == 32
-  return pde + L1_INDEX(va);
-#else
-  if (lvl == 1)
-    return pde + L1_INDEX(va);
-  return pde + L2_INDEX(va);
+static __no_profile inline size_t pde_index(int lvl, vaddr_t va) {
+  switch (lvl) {
+    case 0:
+      return L0_INDEX(va);
+    case 1:
+      return L1_INDEX(va);
+#if __riscv_xlen == 64
+    case 2:
+      return L2_INDEX(va);
+#endif
+  }
+  panic("Invalid level: %d", lvl);
+}
+
+static __no_profile inline bool pde_valid_index(int lvl, size_t index) {
+  switch (lvl) {
+    case 0:
+    case 1:
+#if __riscv_xlen == 64
+    case 2:
 #endif
+      return index < Ln_ENTRIES;
+  }
+  panic("Invalid level: %d", lvl);
+}
+
+static __no_profile inline size_t pde_size(int lvl) {
+  switch (lvl) {
+    case 0:
+      return L0_SIZE;
+    case 1:
+      return L1_SIZE;
+#if __riscv_xlen == 64
+    case 2:
+      return L2_SIZE;
+#endif
+  }
+  panic("Invalid level: %d", lvl);
 }
 
 /*

include/sys/_pmap.h

@@ -70,6 +70,15 @@ paddr_t pde_alloc(pmap_t *pmap);
 pde_t pde_make(int lvl, paddr_t pa);
 static inline pde_t *pde_ptr(paddr_t pd, int lvl, vaddr_t va);
 
+static __no_profile inline pde_t *pde_ptr_idx(paddr_t pd_pa, size_t index) {
+  pde_t *pde = phys_to_dmap(pd_pa);
+  return pde + index;
+}
+
+static __no_profile inline pde_t *pde_ptr(paddr_t pd_pa, int lvl, vaddr_t va) {
+  return pde_ptr_idx(pd_pa, pde_index(lvl, va));
+}
+
 /*
  * Page table.
  */

sys/gen/pmap.c

@@ -315,6 +315,35 @@ bool pmap_extract(pmap_t *pmap, vaddr_t va, paddr_t *pap) {
   return pmap_extract_nolock(pmap, va, pap);
 }
 
+static inline size_t next_vaddr(int lvl) {
+  if (lvl + 1 < PAGE_TABLE_DEPTH)
+    return pde_size(lvl + 1);
+  return PAGESIZE;
+}
+
+static void pmap_protect_walk(int lvl, paddr_t pd_pa, vaddr_t va, vaddr_t end,
+                              vm_prot_t prot) {
+
+  for (size_t pde_i = pde_index(lvl, va);
+       pde_valid_index(lvl, pde_i) && va < end;
+       pde_i++, va += next_vaddr(lvl)) {
+
+    pde_t *pdep = pde_ptr_idx(pd_pa, pde_i);
+    if (!pde_valid_p(pdep))
+      continue;
+
+    paddr_t pa = pte_frame((pte_t)*pdep);
+
+    if (lvl + 1 == PAGE_TABLE_DEPTH) {
+      pte_t *ptep = (pte_t *)pdep;
+      pte_t pte = pte_protect(*ptep, prot);
+      *ptep = pte;
+    } else {
+      pmap_protect_walk(lvl + 1, pa, va, end, prot);
+    }
+  }
+}
+
 void pmap_protect(pmap_t *pmap, vaddr_t start, vaddr_t end, vm_prot_t prot) {
   assert(pmap != pmap_kernel());
   assert(page_aligned_p(start) && page_aligned_p(end));
@@ -324,13 +353,8 @@ void pmap_protect(pmap_t *pmap, vaddr_t start, vaddr_t end, vm_prot_t prot) {
        end);
 
   WITH_MTX_LOCK (&pmap->mtx) {
-    for (vaddr_t va = start; va < end; va += PAGESIZE) {
-      pte_t *ptep = pmap_lookup_pte(pmap, va);
-      if (!pte_valid_p(ptep))
-        continue;
-      pte_t pte = pte_protect(*ptep, prot);
-      pmap_write_pte(pmap, ptep, pte, va);
-    }
+    pmap_protect_walk(0, pmap->pde, start, end, prot);
+    tlb_invalidate_asid(pmap->asid);
   }
 }