#include <stdio.h>
#include <stdlib.h>
#include <sys/mman.h>
#include <string.h>
#include <errno.h>
#include <signal.h>
#include <sys/time.h>

#define MAP_PERM_1              (PROT_WRITE | PROT_READ | PROT_EXEC)
#define MAP_PERM_2              (PROT_WRITE | PROT_READ)

#define MAP_FLAGS               (MAP_ANONYMOUS | MAP_FIXED | MAP_PRIVATE)

#define PG_LEN                  4096
#define align_pg_32(addr)       (addr & 0xFFFFF000)
#define num_pg_in_range(start, end)     ((end - start + 1) >> 12)

inline void __force_pgtbl_alloc(unsigned int start)
{
        volatile int *s = (int *) start;
        *s = *s;
}

int __map_a_page_at(unsigned int start, int whichperm)
{
        int perm = whichperm ? MAP_PERM_1 : MAP_PERM_2;

        if(MAP_FAILED == mmap((void *)start, PG_LEN, perm, MAP_FLAGS, 0, 0)){
                fprintf(stderr,
                        "mmap failed at 0x%x: %s.n",
                        start, strerror(errno));
                return 0;
        }

        return 1;
}

int __mlock_page(unsigned int addr)
{
        if (mlock((void *)addr, (size_t)PG_LEN) < 0){
                fprintf(stderr,
                        "mlock failed on page: 0x%x: %s.n",
                        addr, strerror(errno));
                return 0;
        }

        return 1;
}

void sigint_handler(int p)
{
        struct timeval start = {0 ,0}, end = {0, 0}, diff = {0, 0};
        gettimeofday(&start, NULL);
        munlockall();
        gettimeofday(&end, NULL);
        timersub(&end, &start, &diff);

        printf("Munlock'd entire VM in %u secs %u usecs.n",
                diff.tv_sec, diff.tv_usec);

        exit(0);
}

int make_vma_map(unsigned int start, unsigned int end)
{
        int num_pg = num_pg_in_range(start, end);

        if (end < start){
                fprintf(stderr,
                        "Bad range: start: 0x%x end: 0x%x.n",
                        start, end);
                return 0;
        }

        for (; num_pg; num_pg --, start += PG_LEN){
                if (__map_a_page_at(start, num_pg % 2) && __mlock_page(start))
                        __force_pgtbl_alloc(start);
                else
                        return 0;
        }

        return 1;
}

void display_banner()
{
        printf("-----------------------------------------n");
        printf("Virtual memory allocator. Ctrl+C to exit.n");
        printf("-----------------------------------------n");
}

int main()
{
        unsigned int vma_start, vma_end, input = 0;
        int start_end = 0; // 0: start; 1: end;

        display_banner();

        // Bind SIGINT handler.
        signal(SIGINT, sigint_handler);

        while (1){
                if (!start_end)
                        printf("start:t");
                else
                        printf("end:t");

                scanf("%i", &input);

                if (start_end){
                        vma_end   = align_pg_32(input);
                        make_vma_map(vma_start, vma_end);
                }
                else{
                        vma_start = align_pg_32(input);
                }
                start_end = !start_end;
        }

        return 0;
}