Lab5: Copy-on-Write Fork for xv6

riscv.h

#define PTE_V (1L << 0) // valid
#define PTE_R (1L << 1)
#define PTE_W (1L << 2)
#define PTE_X (1L << 3)
#define PTE_U (1L << 4) // user can access
#define PTE_RSW (1L << 8)

kalloc.c

#define INDEX(pa) (((uint64)pa) / PGSIZE)

int ref_cnt[PHYSTOP / PGSIZE];

void addref(uint64 pa)
{
	acquire(&kmem.lock);
	ref_cnt[INDEX(pa)]++;
	release(&kmem.lock);
}

int getref(uint64 pa)
{
	return ref_cnt[INDEX(pa)];
}

void * kalloc(void)
{
	struct run *r;

	acquire(&kmem.lock);
	r = kmem.freelist;
	if (r)
		kmem.freelist = r->next;
	release(&kmem.lock);

	if (r)
	{
		memset((char *)r, 5, PGSIZE); // fill with junk
		ref_cnt[INDEX(r)] = 1;
	}

	return (void *)r;
}

void kfree(void *pa)
{
	struct run *r;

	if (((uint64)pa % PGSIZE) != 0 || (char *)pa < end || (uint64)pa >= PHYSTOP)
		panic("kfree");

	acquire(&kmem.lock);
  	int remain = --ref_cnt[INDEX(pa)];
  	release(&kmem.lock);

	if (remain > 0)
	{
		return;
	}
	
	memset(pa, 1, PGSIZE);

	r = (struct run *)pa;

	acquire(&kmem.lock);
	r->next = kmem.freelist;
	kmem.freelist = r;
	release(&kmem.lock);
}

void freerange(void *pa_start, void *pa_end)
{
	char *p;
	p = (char *)PGROUNDUP((uint64)pa_start);
	for (; p + PGSIZE <= (char *)pa_end; p += PGSIZE)
	{
		ref_cnt[INDEX(p)] = 1;
		kfree(p);
	}
}

vm.c 

int uvmcopy(pagetable_t old, pagetable_t new, uint64 sz)
{
	pte_t *pte;
	uint64 pa, i;
	uint flags;

	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");
		pa = PTE2PA(*pte);
		if (*pte & PTE_W)
		{
			*pte &= ~PTE_W;
			*pte |= PTE_RSW;
		}
		flags = PTE_FLAGS(*pte);
		if (mappages(new, i, PGSIZE, pa, flags) != 0)
		{
			goto err;
		}
		addref(pa);
	}
	return 0;

err:
	uvmunmap(new, 0, i / PGSIZE, 1);
	return -1;
}

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;
		pte = walk(pagetable, va0, 0);
		if (pte == 0 || (*pte & PTE_V) == 0 || (*pte & PTE_U) == 0)
			return -1;
		if ((*pte & PTE_W) == 0)
		{
			if (cowalloc(pagetable, va0) < 0)
			{
				return -1;
			}
		}
		pa0 = PTE2PA(*pte);
		n = PGSIZE - (dstva - va0);
		if (n > len)
			n = len;
		memmove((void *)(pa0 + (dstva - va0)), src, n);

		len -= n;
		src += n;
		dstva = va0 + PGSIZE;
	}
	return 0;
}

int cowalloc(pagetable_t pagetable, uint64 va)
{
	char* mem;
	pte_t* pte;
	uint64 pa, flags;
	int ref;

	if (va >= MAXVA)
	{
		return -1;
	}

	pte = walk(pagetable, va, 0);
	if (pte == 0 || (*pte & PTE_V) == 0 || (*pte & PTE_U) == 0)
	{
		return -1;
	}
	if ((*pte & PTE_W) == 0 && (*pte & PTE_RSW) == 0)
	{
		return -1;
	}
	
	va = PGROUNDDOWN(va);
	pa = PTE2PA(*pte);
	*pte |= PTE_W;
	*pte &= ~PTE_RSW;
	flags = PTE_FLAGS(*pte);

	ref = getref(pa);

	if (ref == 1)
	{
		return 0;
	}
	else if (ref > 1)
	{
		if ((mem = kalloc()) == 0)
		{
			return -1;
		}

		memmove(mem, (char *)pa, PGSIZE);
		uvmunmap(pagetable, va, 1, 1);
		if (mappages(pagetable, va, PGSIZE, (uint64)mem, flags) != 0)
		{
			kfree(mem);
			return -1;
		}
		return 0;
	}
	else
	{
		return -1;
	}
}

 trap.c

void usertrap(void)
{
    
    ......

	else if (r_scause() == 15)
	{
		if (cowalloc(p->pagetable, r_stval()) < 0)
		{
			p->killed = 1;
		}
	}
	
	......
}