Reference count page tables

kernel/mm/bba.h

@@ -38,7 +38,7 @@ SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 /* Bootstrap Bitmap Allocator */
 #define BBA_BLOCK_SIZE (4096)
-#define BBA_SIZE      (BBA_BLOCK_SIZE * 128)
+#define BBA_SIZE      (BBA_BLOCK_SIZE * 256)
 #define BBA_BM_SIZE   (BBA_BLOCK_SIZE)
 
 struct bba {

kernel/platform/i386_pc/sys/mm.c

@@ -52,6 +52,7 @@ static struct page_dir *mm_alloc_pd(void) {
   memset(pd, 0, sizeof(*pd));
   sl_init(&pd->sl, "pd");
   pd->pd = (volatile uint32_t *)bba_alloc();
+  memset(pd->pt_refcount, 0, sizeof(pd->pt_refcount));
   return pd;
 }
 
@@ -60,6 +61,7 @@ static struct page_dir *mm_alloc_existing_pd(uptr_t ptr) {
   memset(pd, 0, sizeof(*pd));
   sl_init(&pd->sl, "pd");
   pd->pd = (volatile uint32_t *)ptr;
+  memset(pd->pt_refcount, 0, sizeof(pd->pt_refcount));
   return pd;
 }
 
@@ -85,9 +87,13 @@ void mm_map_page(struct page_dir *pd, uptr_t vaddr, uptr_t paddr, uint32_t flags
       pt[i] = 0;
     uptr_t phys_pt = (uptr_t)pt - VIRT_BASE;
     pd->pd[pd_idx] = phys_pt | PF_PRESENT | PF_WRITABLE;
+    pd->pt_refcount[pd_idx] = 1;
   } else {
     uptr_t pt_phys = pd->pd[pd_idx] & ~(PAGE_SIZE - 1);
     pt = (volatile uint32_t *)(pt_phys + VIRT_BASE);
+
+    if (!(pt[pt_idx] & PF_PRESENT))
+      pd->pt_refcount[pd_idx]++;
   }
 
   pt[pt_idx] = (paddr & ~(PAGE_SIZE - 1)) | (flags & ~PF_LOCK);
@@ -112,8 +118,18 @@ void mm_unmap_page(struct page_dir *pd, uptr_t vaddr, uint32_t flags) {
     pt = (volatile uint32_t *)(pt_phys + VIRT_BASE);
   }
 
+  if (pt[pt_idx] & PF_PRESENT) {
     pt[pt_idx] &= ~PF_PRESENT;
 
+    if (pd->pt_refcount[pd_idx] > 0)
+      pd->pt_refcount[pd_idx]--;
+
+    if (pd->pt_refcount[pd_idx] == 0) {
+      bba_free((uptr_t)pt);
+      pd->pd[pd_idx] = 0;
+    }
+  }
+  
   __asm__ volatile ("invlpg (%0)" :: "r"(vaddr) : "memory");
 
 done:

kernel/platform/i386_pc/sys/mm.h

@@ -49,10 +49,10 @@ SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
                                  it would be inefficient to constantly (un)lock. */
 
 #define KERNEL_HEAP_START       0xF0000000
-#define MEM_PROBE_HEAP_START 0xE0000000
 
 struct page_dir {
   volatile uint32_t *pd;
+  uint16_t pt_refcount[1024];
   struct spinlock sl;
 };