#include <linux/mm.h>
#include <linux/mm_types.h>
#include <linux/sched.h>
#include <asm/pgalloc.h>
#include <asm/tlbflush.h>

// vaddr_h: vaddr above 4G
long sys_set_pte(unsigned long vaddr_h)
{
    unsigned long vaddr = 0x10000000 | (vaddr_h & 0xfff); // make vaddr's pte point to vaddr_h's phys page
    // vaddr and vaddr_h shares the same pgd and pud, need to creat pmd and page table
    pgd_t *pgd;
    pud_t *pud;
    pmd_t *pmd, pmd_n; // pmd_n is for vaddr
    pte_t *pte, *pte_h;
    struct page *page;
    pgprot_t pgprot;
    
    // get vaddr_h's physical page
    pgd = pgd_offset(current->mm, vaddr_h);
    pud = pud_offset(pgd, vaddr_h);
    pmd = pmd_offset(pud, vaddr_h);
    pgprot = __pgprot(pmd_val(*pmd) & 0xfff); // we copy vaddr_h's pmd permission to vaddr's pmd
    pte_h = pte_offset_map(pmd, vaddr_h);
    
    // mapping vaddr_h above 4G to vaddr below 4G alloc page entry
    pgd = pgd_offset(current->mm, vaddr);
    if (pgd_none(*pgd)) {
        printk("pgd entry not found, alloc new pud and set pgd entry\n");
        pgd = pgd_alloc(current->mm);
    }

    pud = pud_offset(pgd, vaddr);
    if(pud_none(*pud)) {
        printk("pud entry not found, alloc new pmd and set pud entry\n");
        pud = pud_alloc(current->mm, pgd, vaddr);
    }
    
    pmd = pmd_offset(pud, vaddr);
    if(pmd_none(*pmd)) {
        printk("pmd entry not found, alloc new pte and set pmd entry\n");
        page = pte_alloc_one(current->mm, vaddr); // allocate pte, i.e., page table
        pmd_n = mk_pmd(page, pgprot); // make a new pmd entry which ponits to page with pgprot permission
        set_pmd(pmd, pmd_n); // replace old pmd entry (pmd) with a new one (pmd_n)
    }

    // pte = page table entry
    pte = pte_offset_map(pmd, vaddr);

    // replace vaddr page table entry (pte) with vaddr_h's one (*pte_h)
    set_pte(pte, *pte_h);
}