extern crate std;
use criterion::{black_box, criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
extern crate iai;

use core::borrow::Borrow;
use core::cmp::{Ord, Ordering};
use core::ptr;
use core::slice;

fn std_binary<K>(mut pkey: *const K, len: usize, key: &K) -> Result<usize, usize>
where
    K: Ord,
{
    let slice = unsafe { slice::from_raw_parts(pkey, len) };
    slice.binary_search(key)
}

fn slice_linear<K, Q>(mut pkey: *const K, len: usize, key: &Q) -> Result<usize, usize>
where
    K: Borrow<Q> + Ord,
    Q: Ord,
{
    let slice = unsafe { slice::from_raw_parts(pkey, len) };
    for item in slice {
        let ret = key.cmp(item.borrow());
        if ret == Ordering::Greater {
            continue;
        } else if ret == Ordering::Less {
            unsafe { return Err((item as *const K).offset_from(pkey) as usize) };
        } else {
            unsafe { return Ok((item as *const K).offset_from(pkey) as usize) };
        }
    }
    Err(len)
}

fn idx_linear<K, Q>(mut pkey: *const K, len: usize, key: &Q) -> Result<usize, usize>
where
    K: Borrow<Q> + Ord,
    Q: Ord,
{
    let slice = unsafe { slice::from_raw_parts(pkey, len) };
    for idx in 0..len {
        let item = (&slice[idx]).borrow();
        let ret = key.cmp(item);
        if ret == Ordering::Greater {
            continue;
        } else if ret == Ordering::Less {
            return Err(idx);
        } else {
            return Ok(idx);
        }
    }
    Err(len)
}

fn slice_binary<K, Q>(mut pkey: *const K, mut len: usize, key: &Q) -> Result<usize, usize>
where
    K: Borrow<Q> + Ord,
    Q: Ord,
{
    let slice = unsafe { slice::from_raw_parts(pkey, len) };
    let mut lo = 0;
    while len > 0 {
        let mid = lo + len / 2;
        //let item = unsafe { &*pkey.add(mid) }.borrow();
        let item = unsafe { slice.get_unchecked(mid).borrow() };
        let ret = key.cmp(item);
        if ret == Ordering::Greater {
            len -= mid - lo + 1;
            lo = mid + 1;
        } else if ret == Ordering::Less {
            len = mid - lo;
        } else {
            return Ok(mid);
        }
    }
    Err(lo)
}

fn idx_binary<K, Q>(mut pkey: *const K, mut len: usize, key: &Q) -> Result<usize, usize>
where
    K: Borrow<Q> + Ord,
    Q: Ord,
{
    let slice = unsafe { slice::from_raw_parts(pkey, len) };
    let mut lo = 0;
    while len > 0 {
        let mid = lo + len / 2;
        let item = (&slice[mid]).borrow();
        let ret = key.cmp(item);
        if ret == Ordering::Greater {
            len -= mid - lo + 1;
            lo = mid + 1;
        } else if ret == Ordering::Less {
            len = mid - lo;
        } else {
            return Ok(mid);
        }
    }
    Err(lo)
}

fn ptr_linear<K, Q>(mut pkey: *const K, len: usize, key: &Q) -> Result<usize, usize>
where
    K: Borrow<Q> + Ord,
    Q: Ord,
{
    for idx in 0..len {
        let item = unsafe { (&*pkey).borrow() };
        let ret = key.cmp(item);
        if ret == Ordering::Greater {
            unsafe { pkey = pkey.add(1) };
            continue;
        } else if ret == Ordering::Less {
            return Err(idx);
        } else {
            return Ok(idx);
        }
    }
    Err(len)
}

fn ptr_binary<K, Q>(mut pkey: *const K, mut len: usize, key: &Q) -> Result<usize, usize>
where
    K: Borrow<Q> + Ord,
    Q: Ord,
{
    let mut lo = 0;
    while len > 0 {
        let mid = lo + len / 2;
        let item = unsafe { &*pkey.add(mid) }.borrow();
        let ret = key.cmp(item);
        if ret == Ordering::Greater {
            len -= mid - lo + 1;
            lo = mid + 1;
        } else if ret == Ordering::Less {
            len = mid - lo;
        } else {
            return Ok(mid);
        }
    }
    Err(lo)
}

struct SearchCtx {
    data: Vec<usize>,
    keys: Vec<usize>,
}

static mut CTX: Option<SearchCtx> = None;

impl SearchCtx {
    fn alloc(cnt: usize) -> Self {
        let mut vec = Self {
            data: vec![0; cnt],
            keys: vec![0; 1000],
        };
        vec.keys.iter_mut().for_each(|item| {
            *item = rand::random::<usize>();
        });
        vec.data.iter_mut().for_each(|item| {
            *item = rand::random::<usize>();
        });
        vec.data.sort();
        vec
    }
    fn new(cnt: usize) -> Self {
        unsafe {
            if let Some(ref ctx) = CTX {
                if ctx.data.len() != cnt {
                    CTX = Some(Self::alloc(cnt));
                }
            } else {
                CTX = Some(Self::alloc(cnt));
            }

            let ctx = CTX.as_ref().unwrap();
            Self {
                data: ctx.data.clone(),
                keys: ctx.data.clone(),
            }
        }
    }

    fn search<F>(&self, f: F)
    where
        F: Fn(*const usize, usize, &usize) -> Result<usize, usize>,
    {
        for key in &self.keys {
            f(self.data.as_ptr(), self.data.len(), &key);
        }
    }
}

fn search_binary(c: &mut Criterion) {
    let mut group = c.benchmark_group("binary");
    let cnts = [10, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8096];
    for i in 0..cnts.len() {
        group.bench_function(BenchmarkId::new("std_binary", i), |b| {
            b.iter_batched_ref(
                || -> SearchCtx { SearchCtx::new(cnts[i]) },
                |ctx| ctx.search(std_binary),
                BatchSize::SmallInput,
            )
        });
        group.bench_function(BenchmarkId::new("slice_binary", i), |b| {
            b.iter_batched_ref(
                || -> SearchCtx { SearchCtx::new(cnts[i]) },
                |ctx| ctx.search(slice_binary),
                BatchSize::SmallInput,
            )
        });
        group.bench_function(BenchmarkId::new("idx_binary", i), |b| {
            b.iter_batched_ref(
                || -> SearchCtx { SearchCtx::new(cnts[i]) },
                |ctx| ctx.search(idx_binary),
                BatchSize::SmallInput,
            )
        });
        group.bench_function(BenchmarkId::new("ptr_binary", i), |b| {
            b.iter_batched_ref(
                || -> SearchCtx { SearchCtx::new(cnts[i]) },
                |ctx| ctx.search(ptr_binary),
                BatchSize::SmallInput,
            )
        });
    }
    group.finish();
}

fn search_linear(c: &mut Criterion) {
    let mut group = c.benchmark_group("linear");
    let cnts = [10, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8096];
    for i in 0..cnts.len() {
        group.bench_function(BenchmarkId::new("slice_linear", i), |b| {
            b.iter_batched_ref(
                || -> SearchCtx { SearchCtx::new(cnts[i]) },
                |ctx| ctx.search(slice_linear),
                BatchSize::SmallInput,
            )
        });
        group.bench_function(BenchmarkId::new("ptr_linear", i), |b| {
            b.iter_batched_ref(
                || -> SearchCtx { SearchCtx::new(cnts[i]) },
                |ctx| ctx.search(ptr_linear),
                BatchSize::SmallInput,
            )
        });
        group.bench_function(BenchmarkId::new("idx_linear", i), |b| {
            b.iter_batched_ref(
                || -> SearchCtx { SearchCtx::new(cnts[i]) },
                |ctx| ctx.search(idx_linear),
                BatchSize::SmallInput,
            )
        });
    }
    group.finish();
}

criterion_group!(benches, search_linear, search_binary);
criterion_main!(benches);
