rust self-tracing benchmark using DebugCtl.BTF

rust self-tracing benchmark using DebugCtl.BTF.rs

#![feature(const_fn)]

// rust self-tracing benchmark using DebugCtl.BTF

extern crate winapi;
extern crate kernel32;
extern crate libc;

use std::time::Instant;

use libc::{c_long}; 
pub type LONG = c_long;

use std::vec::Vec;

use winapi::ntstatus::STATUS_SINGLE_STEP;
use winapi::winnt::EXCEPTION_POINTERS;
use winapi::winnt::CONTEXT;

use kernel32::AddVectoredExceptionHandler;

extern "system" {
    pub fn RaiseException(dwExceptionCode: u32,
                          dwExceptionFlags: u32,
                          nNumberOfArguments: u32,
                          lpArguments: *const u64);
}

const EXCEPTION_CONTINUE_SEARCH: i32 = 0;
const EXCEPTION_CONTINUE_EXECUTION: i32 = -1;
const DEBUGCTL_BTF_FLAG: u64  = 0x300; 
const DEBUGCTL_TRAP_FLAG: u32 = 0x100; 

extern "system" fn trace_callback(ex: *mut EXCEPTION_POINTERS) -> LONG
{
    unsafe { 
        let ctx: *mut CONTEXT = (*ex).ContextRecord;
        let excr = &(*(*ex).ExceptionRecord);

        match excr.ExceptionCode as i32 {
           STATUS_SINGLE_STEP => { 
               (*blocks).push((*ctx).Rip);
               (*ctx).Dr7 |= DEBUGCTL_BTF_FLAG;
               (*ctx).EFlags |= DEBUGCTL_TRAP_FLAG;
               return EXCEPTION_CONTINUE_EXECUTION;
            }
            _ => { return EXCEPTION_CONTINUE_SEARCH; }         
        }
    }
}


// intentionally slow is_prime functoin
fn is_prime(n: u32) -> bool {
    for i in 2..n
    {
        if n % i == 0
        {
            return false;
        }
    }
    return true;
}

// ptr to buffer for exception handler to log basic blocks
static mut blocks : *mut Vec<u64> = std::ptr::null_mut();

fn main() {
    let mut b : Vec<u64> = Vec::with_capacity(1000000);
    
    unsafe {
    blocks = &mut b;
    AddVectoredExceptionHandler(1, Some(trace_callback));
    }

    let mut t1 = Instant::now();
    for i in 1..1000
    {
        if is_prime(i) {
            //println!("{0} is prime", i);
        }
    }
    let mut t2 = Instant::now();
    println!("native is_prime {:?}", t2.duration_since(t1));


    // enable block stepping and do it again 
    println!("Enabling block trace..");
    unsafe{
    RaiseException(STATUS_SINGLE_STEP as u32, 0, 0, std::ptr::null());
    }

    t1 = Instant::now();
    for i in 1..1000
    {
        if is_prime(i) {
            //println!("{0} is prime", i);
        }
    }
    t2 = Instant::now();
    println!("traced is_prime {:?}", t2.duration_since(t1));

    println!("Basic blocks:");
    for i in 0..b.len() {
        println!("    {0:x}", b[i]);
    }
    println!("Traced {0} blocks", b.len());

}