///////////////////////////////////////////////////////////////////////////////
//////////////////// Section 1: stuff to make this valid C ////////////////////

#include <stdlib.h>

#define NULL ((void *)0)
typedef unsigned long ulong;
typedef int bool;

void *get_new_physical_page() {
	// FIX ME: this method should guarantee that the address returned is a
	// multiple of page_size, but this implementation does not guarantee that.
	return malloc(page_size);
}

//////////////////// Section 1: stuff to make this valid C ////////////////////
///////////////////////////////////////////////////////////////////////////////


///////////////////////////////////////////////////////////////////////////////
/////////////////// Section 2: parameters MMU desginers pick //////////////////

#define po_bits   12
#define pa_bits   52
#define levels     4

void *ptbr; // the page table base register - a kernel-mode register

/////////////////// Section 2: parameters MMU desginers pick //////////////////
///////////////////////////////////////////////////////////////////////////////


///////////////////////////////////////////////////////////////////////////////
////////////////////// Section 3: physical address stuff //////////////////////

// a page contains 2**po_bits bytes of data
#define page_size (1<<po_bits)
// the part of an address that is not the offset is/are the page number(s)
#define ppn_bits  (pa_bits-po_bits)

/** a physical page number is the high-order bits of a physical address */
void *pa_from_ppn(ulong ppn) {
	return (void *)(ppn<<po_bits);
}
/** a physical page number is the high-order bits of a physical address */
ulong ppn_from_pa(void *pa) {
	return (((ulong)pa)>>po_bits) & ((1L<<pa_bits)-1);
}

////////////////////// Section 3: physical address stuff //////////////////////
///////////////////////////////////////////////////////////////////////////////



///////////////////////////////////////////////////////////////////////////////
/////////////////////// Section 4: virtual address stuff //////////////////////

// We want a page table entry to take a power-of-2 total bytes
// we can enforce that by picking a big-enough power of 2...
#define log_pte_bytes 3
#define pte_bytes     (1<<log_pte_bytes)
// ...and by using "packed struct" syntax to set bit width for each field
typedef struct PTE_t {
	ulong ppn       : ppn_bits;
	bool  allocated : 1;
	// (in real code we'd have other protection bits here...)
	ulong _unused   : ((pte_bytes*8) - (ppn_bits+1)); // make PTE right length
} PTE;

// a page table is an array of page table entries that fits in one page,
// and a single level of the virtual page number is an index into that array
#define one_vpn_bits  (po_bits - log_pte_bytes)
#define pte_per_pt    (1<<one_vpn_bits) // == (page_size/sizeof(PTE))
typedef struct PT_t {
	PTE entries[pte_per_pt];
} PT; 

// we never use the following values, but here is how to compute them:
#define all_vpn_bits  (levels * one_vpn_bits)
#define va_bits       (all_vpn_bits + po_bits)

/////////////////////// Section 4: virtual address stuff //////////////////////
///////////////////////////////////////////////////////////////////////////////


///////////////////////////////////////////////////////////////////////////////
//////////////// Section 5: address translation and allocation ////////////////

/**
 * Turns a virtual address into a physical address,
 * or returns NULL if there is a page fault.
 */
void *pa_from_va(ulong va) {
	ulong po = va & (page_size-1);
	ulong vpns = va >> po_bits;

	void *page = ptbr;
	for(int i=0; i<levels; i+=1) {
		ulong vpn = (vpns >> ((3-i)*one_vpn_bits))&(pte_per_pt-1);    // get next level of VPN
		
		PTE pte = ((PT *)page)->entries[vpn];
		
		if (!pte.allocated) return NULL;    // page fault
		
		page = pa_from_ppn(pte.ppn);        // next level of tree
	}
	return page + po;
}

/**
 * If the given virtual address is already allocated, does nothing.
 * Otherwise, allocates the page containing that virtual address
 * and all intermediate page tables.
 */
void allocate(ulong va) {
	 // the page offset does not impact allocation
	ulong vpns = va >> po_bits;

	void *page = ptbr;
	for(int i=0; i<levels; i+=1) {
		ulong vpn = (vpns >> ((3-i)*one_vpn_bits))&(pte_per_pt-1);    // get next level of VPN
		
		PTE pte = ((PT *)page)->entries[vpn];
		
		if (!pte.allocated) {
			// allocate a new page
			void *new_page = get_new_physical_page();
			// ensure all its entries are marked "unallocated"
			for(int j=0; j<pte_per_pt; j+=1)
				((PT *)new_page)->entries[j].allocated = 0;
			// update the page table entry
			pte.ppn = ppn_from_pa(new_page);
			pte.allocated = 1;
			((PT *)page)->entries[vpn] = pte;
		}
		
		page = pa_from_ppn(pte.ppn);        // next level of tree
	}
}

//////////////// Section 5: address translation and allocation ////////////////
///////////////////////////////////////////////////////////////////////////////


void main(void) {
	levels = 4;
	ptbr = calloc(page_size, 0);
	
	printf("Page fault: %p\n", pa_from_va(123456));
	allocate(123456);
	printf("Page hit:   %p\n", pa_from_va(123456));
	*(int *)pa_from_va(123456) = 0x3330;
	printf("Value in memory: %x\n", *(int *)pa_from_va(123456));
	printf("Different part of page: %p\n", pa_from_va(123450));
	printf("Different page: %p\n", pa_from_va(120000));
}