extern crate time;

use time::precise_time_ns;

use std::mem;
use std::process::exit;

const PAGE_SIZE: usize = 4096;

fn main() {
    let args: Vec<String> = std::env::args().collect();

    match args.len() {
        0 | 1 => pow2_series(1_000_000, 16),
        3 => {
            match parse_args(&args[1], &args[2]) {
                Ok(time) => println!("{} ns", time),
                Err(err) => println!("{}", err),
            }
        }
        _ => {
            println!("Wrong number of arguments");
            exit(1);
        }
    }
}

fn parse_args(page_count: &str, trail_count: &str) -> Result<f64, &'static str> {
    let page_count = page_count
        .parse::<usize>()
        .map_err(|_| "Illegal page_count")?;
    let trail_count = trail_count
        .parse::<u32>()
        .map_err(|_| "Illegal trail_count")?;

    Ok(tlb_probe(page_count, trail_count))
}

/// Prints the results of multiple tlb_probe.
/// The series has a given length and page_count gets incremented in powers of 2.
/// The given trail_size is normalized for each run to get a consistent executing time.
fn pow2_series(trail_count: u32, length: u32) {
    for page_count in (0..length).map(|x| 2u64.pow(x)) {
        let r = tlb_probe(page_count as usize, trail_count / page_count as u32);
        println!("{} = {:?}", page_count, r);
    }
}

/// Probes the tlb using the two parameters.
/// Returns the time taken per operation in ns.
fn tlb_probe(page_count: usize, trail_count: u32) -> f64 {
    let jump = PAGE_SIZE / mem::size_of::<i32>();

    let mut a = vec![0i32; page_count * jump];

    // warm up
    for i in a.iter_mut() {
        *i += 1;
    }

    let before = precise_time_ns();

    // the hot loop
    for _ in 0..trail_count {
        let mut i = 0;

        while i < a.len() {
            a[i] += 1;

            i += jump;
        }
    }

    let after = precise_time_ns();

    // time calculations
    let d = after - before;

    d as f64 / (trail_count * page_count as u32) as f64
}