lab6.txt

diff --git a/kernel/defs.h b/kernel/defs.h
index 4b9bbc0..9998537 100644
--- a/kernel/defs.h
+++ b/kernel/defs.h
@@ -63,6 +63,11 @@ void            ramdiskrw(struct buf*);
 void*           kalloc(void);
 void            kfree(void *);
 void            kinit(void);
+int             get_kmem_ref(void *pa);
+void            add_kmem_ref(void *pa);
+void            acquire_ref_lock();
+void            release_ref_lock();
+
 
 // log.c
 void            initlog(int, struct superblock*);
diff --git a/kernel/kalloc.c b/kernel/kalloc.c
index fa6a0ac..e826444 100644
--- a/kernel/kalloc.c
+++ b/kernel/kalloc.c
@@ -18,16 +18,28 @@ struct run {
   struct run *next;
 };
 
+#define INDEX(pa) (((char*)pa - (char*)PGROUNDUP((uint64)end)) >> 12)
+
 struct {
   struct spinlock lock;
   struct run *freelist;
+  // add ref
+  struct spinlock ref_lock;
+  uint *ref_count;
 } kmem;
 
 void
 kinit()
 {
   initlock(&kmem.lock, "kmem");
-  freerange(end, (void*)PHYSTOP);
+  initlock(&kmem.ref_lock, "ref");
+
+  //total physical pages
+  uint64 physical_pages = ((PHYSTOP - (uint64)end) >> 12) + 1;
+  physical_pages = ((physical_pages * sizeof(uint)) >> 12) + 1;
+  kmem.ref_count = (uint*) end;
+  uint64 offset = physical_pages << 12;
+  freerange(end + offset, (void*)PHYSTOP);
 }
 
 void
@@ -35,8 +47,10 @@ freerange(void *pa_start, void *pa_end)
 {
   char *p;
   p = (char*)PGROUNDUP((uint64)pa_start);
-  for(; p + PGSIZE <= (char*)pa_end; p += PGSIZE)
+  for(; p + PGSIZE <= (char*)pa_end; p += PGSIZE){
+    kmem.ref_count[INDEX((void*)p)] = 1;
     kfree(p);
+  }
 }
 
 // Free the page of physical memory pointed at by v,
@@ -51,6 +65,14 @@ kfree(void *pa)
   if(((uint64)pa % PGSIZE) != 0 || (char*)pa < end || (uint64)pa >= PHYSTOP)
     panic("kfree");
 
+  // check if the number reaches 0
+  acquire(&kmem.lock);
+  if (--kmem.ref_count[INDEX(pa)]){
+    release(&kmem.lock);
+    return;
+  }
+  release(&kmem.lock);
+
   // Fill with junk to catch dangling refs.
   memset(pa, 1, PGSIZE);
 
@@ -72,11 +94,29 @@ kalloc(void)
 
   acquire(&kmem.lock);
   r = kmem.freelist;
-  if(r)
+  if(r){
     kmem.freelist = r->next;
+    kmem.ref_count[INDEX((void*)r)] = 1;
+  }
   release(&kmem.lock);
 
   if(r)
     memset((char*)r, 5, PGSIZE); // fill with junk
   return (void*)r;
 }
+
+int get_kmem_ref(void *pa){
+  return kmem.ref_count[INDEX(pa)];
+}
+ 
+void add_kmem_ref(void *pa){
+  kmem.ref_count[INDEX(pa)]++;
+}
+ 
+void acquire_ref_lock(){
+  acquire(&kmem.ref_lock);
+}
+ 
+void release_ref_lock(){
+  release(&kmem.ref_lock);
+}
\ No newline at end of file
diff --git a/kernel/riscv.h b/kernel/riscv.h
index 0aec003..726d6d4 100644
--- a/kernel/riscv.h
+++ b/kernel/riscv.h
@@ -332,6 +332,8 @@ sfence_vma()
 #define PTE_X (1L << 3)
 #define PTE_U (1L << 4) // 1 -> user can access
 
+#define PTE_RSW (1L << 8) // RSW
+
 // shift a physical address to the right place for a PTE.
 #define PA2PTE(pa) ((((uint64)pa) >> 12) << 10)
 
diff --git a/kernel/trap.c b/kernel/trap.c
index a63249e..10f7580 100644
--- a/kernel/trap.c
+++ b/kernel/trap.c
@@ -6,6 +6,8 @@
 #include "proc.h"
 #include "defs.h"
 
+extern pte_t* walk(pagetable_t, uint64, int);
+
 struct spinlock tickslock;
 uint ticks;
 
@@ -65,7 +67,47 @@ usertrap(void)
     intr_on();
 
     syscall();
-  } else if((which_dev = devintr()) != 0){
+  } 
+  else if (r_scause() == 15) { // write page fault
+    uint64 va = PGROUNDDOWN(r_stval());
+    pte_t *pte;
+    if (va > p->sz || (pte = walk(p->pagetable, va, 0)) == 0){
+      p->killed = 1;
+      goto end;
+    }
+
+    if (((*pte) & PTE_RSW) == 0 || ((*pte) & PTE_V) == 0 || ((*pte) & PTE_U) == 0){
+      p->killed = 1;
+      goto end;
+    }
+
+    uint64 pa = PTE2PA(*pte);
+    acquire_ref_lock();
+    uint ref = get_kmem_ref((void*)pa);
+    if (ref == 1){
+      *pte = ((*pte) & (~PTE_RSW)) | PTE_W;
+    }
+    else {
+      char* mem = kalloc();
+      if (mem == 0){
+        p->killed = 1;
+        release_ref_lock();
+        goto end;
+      }
+
+      memmove(mem, (char*)pa, PGSIZE);
+      uint flag = (PTE_FLAGS(*pte) | PTE_W) & (~PTE_RSW);
+      if (mappages(p->pagetable, va, PGSIZE, (uint64)mem,flag) != 0){
+        kfree(mem);
+        p->killed = 1;
+        release_ref_lock();
+        goto end;
+      }
+      kfree((void*)pa);
+    }
+    release_ref_lock();
+  }
+  else if((which_dev = devintr()) != 0){
     // ok
   } else {
     printf("usertrap(): unexpected scause %p pid=%d\n", r_scause(), p->pid);
@@ -73,6 +115,7 @@ usertrap(void)
     p->killed = 1;
   }
 
+end:
   if(p->killed)
     exit(-1);
 
diff --git a/kernel/vm.c b/kernel/vm.c
index bccb405..dd2d744 100644
--- a/kernel/vm.c
+++ b/kernel/vm.c
@@ -6,6 +6,8 @@
 #include "defs.h"
 #include "fs.h"
 
+extern pte_t* walk(pagetable_t, uint64, int);
+
 /*
  * the kernel's page table.
  */
@@ -156,8 +158,8 @@ mappages(pagetable_t pagetable, uint64 va, uint64 size, uint64 pa, int perm)
   for(;;){
     if((pte = walk(pagetable, a, 1)) == 0)
       return -1;
-    if(*pte & PTE_V)
-      panic("remap");
+    // if(*pte & PTE_V)
+    //   panic("remap");
     *pte = PA2PTE(pa) | perm | PTE_V;
     if(a == last)
       break;
@@ -311,22 +313,25 @@ uvmcopy(pagetable_t old, pagetable_t new, uint64 sz)
   pte_t *pte;
   uint64 pa, i;
   uint flags;
-  char *mem;
 
   for(i = 0; i < sz; i += PGSIZE){
     if((pte = walk(old, i, 0)) == 0)
       panic("uvmcopy: pte should exist");
     if((*pte & PTE_V) == 0)
       panic("uvmcopy: page not present");
+    // set unwrite and set rsw
+    *pte = ((*pte) & (~PTE_W)) | PTE_RSW;
     pa = PTE2PA(*pte);
     flags = PTE_FLAGS(*pte);
-    if((mem = kalloc()) == 0)
-      goto err;
-    memmove(mem, (char*)pa, PGSIZE);
-    if(mappages(new, i, PGSIZE, (uint64)mem, flags) != 0){
-      kfree(mem);
+    // stop copy
+    // if((mem = kalloc()) == 0)
+    //   goto err;
+    // memmove(mem, (char*)pa, PGSIZE);
+    if(mappages(new, i, PGSIZE, pa, flags) != 0){
+      // kfree(mem);
       goto err;
     }
+    add_kmem_ref((void*)pa);
   }
   return 0;
 
@@ -355,9 +360,41 @@ int
 copyout(pagetable_t pagetable, uint64 dstva, char *src, uint64 len)
 {
   uint64 n, va0, pa0;
+  pte_t* pte;
 
   while(len > 0){
     va0 = PGROUNDDOWN(dstva);
+    if(va0 >= MAXVA)  
+      return -1;
+    if((pte = walk(pagetable, va0, 0)) == 0)
+      return -1;
+    if (((*pte & PTE_V) == 0) || ((*pte & PTE_U)) == 0) 
+      return -1;
+    
+    pa0 = PTE2PA(*pte);
+    if (((*pte & PTE_W) == 0) && (*pte & PTE_RSW)){
+      acquire_ref_lock();
+      if (get_kmem_ref((void*)pa0) == 1) {
+          *pte = (*pte | PTE_W) & (~PTE_RSW);
+      }
+      else {
+        char* mem = kalloc();
+        if (mem == 0){
+          release_ref_lock();
+          return -1;
+        }
+        memmove(mem, (char*)pa0, PGSIZE);
+        uint new_flags = (PTE_FLAGS(*pte) | PTE_RSW) & (~PTE_W);
+        if (mappages(pagetable, va0, PGSIZE, (uint64)mem, new_flags) != 0){
+          kfree(mem);
+          release_ref_lock();
+          return -1;
+        }
+        kfree((void*)pa0);
+      }
+      release_ref_lock();
+    }
+    
     pa0 = walkaddr(pagetable, va0);
     if(pa0 == 0)
       return -1;