<commit_before>// Copyright 2015 Keegan McAllister.
// Copyright 2016 Corey Farwell.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// See `LICENSE` in this repository.

use std::io::{self, Read};
use std::{panic, process};

/// Utility that reads a `Vec` of bytes from standard input (stdin)
/// and passes it to `closure`. All panics that occur within
/// `closure` will be treated as aborts. This is done so that
/// AFL considers a panic to be a crash.
///
/// # Examples
///
/// ```no_run
/// extern crate afl;
/// # struct Image;
/// # impl Image {
/// #     fn parse(_: &[u8]) {}
/// # }
///
/// fn main() {
///     afl::read_stdio_bytes(|read| {
///         Image::parse(&read)
///     })
/// }
/// ```
#[deprecated(
    since = "0.3.3",
    note = "This function does not use the `persistent mode` and `defered forkserver mode` and is therefore very slow.
Please use fuzz() or fuzz!() instead."
)]
pub fn read_stdio_bytes<F>(closure: F)
where
    F: Fn(Vec<u8>) + panic::RefUnwindSafe,
{
    let mut input = vec![];
    let result = io::stdin().read_to_end(&mut input);
    if result.is_err() {
        return;
    }
    let was_panic = panic::catch_unwind(|| {
        closure(input);
    });
    if was_panic.is_err() {
        process::abort();
    }
}

/// Utility that reads a `String` from standard input (stdin) and
/// passes it to `closure`. If a `String` cannot be constructed from
/// the data provided by standard input, `closure` will _not_ be
/// called. All panics that occur within `closure` will be treated as
/// aborts. This is done so that AFL considers a panic to be a crash.
///
/// # Examples
///
/// ```no_run
/// extern crate afl;
/// # struct Url;
/// # impl Url {
/// #     fn parse(_: &str) {}
/// # }
///
/// fn main() {
///     afl::read_stdio_string(|string| {
///         Url::parse(&string)
///     })
/// }
/// ```
#[deprecated(
    since = "0.3.3",
    note = "This function does not use the `persistent mode` and `defered forkserver mode` and is therefore very slow.
Please use fuzz() or fuzz!() instead."
)]
pub fn read_stdio_string<F>(closure: F)
where
    F: Fn(String) + panic::RefUnwindSafe,
{
    let mut input = String::new();
    let result = io::stdin().read_to_string(&mut input);
    if result.is_err() {
        return;
    }
    let was_panic = panic::catch_unwind(|| {
        closure(input);
    });
    if was_panic.is_err() {
        process::abort();
    }
}

// those functions are provided by the afl-llvm-rt static library
extern "C" {
    fn __afl_persistent_loop(counter: usize) -> isize;
    fn __afl_manual_init();
}

/// Fuzz a closure by passing it a `&[u8]`
///
/// This slice contains a "random" quantity of "random" data.
///
/// ```rust,no_run
/// # extern crate afl;
/// # use afl::fuzz;
/// # fn main() {
/// fuzz(|data|{
///     if data.len() != 6 {return}
///     if data[0] != b'q' {return}
///     if data[1] != b'w' {return}
///     if data[2] != b'e' {return}
///     if data[3] != b'r' {return}
///     if data[4] != b't' {return}
///     if data[5] != b'y' {return}
///     panic!("BOOM")
/// });
/// # }
/// ```
pub fn fuzz<F>(closure: F)
where
    F: Fn(&[u8]) + std::panic::RefUnwindSafe,
{
    // this marker strings needs to be in the produced executable for
    // afl-fuzz to detect `persistent mode` and `defered mode`
    static PERSIST_MARKER: &'static str = "##SIG_AFL_PERSISTENT##\0";
    static DEFERED_MARKER: &'static str = "##SIG_AFL_DEFER_FORKSRV##\0";

    // we now need a fake instruction to prevent the compiler from optimizing out
    // those marker strings
    unsafe { std::ptr::read_volatile(&PERSIST_MARKER) }; // hack used in https://github.com/bluss/bencher for black_box()
    unsafe { std::ptr::read_volatile(&DEFERED_MARKER) };
    // unsafe { asm!("" : : "r"(&PERSIST_MARKER)) }; // hack used in nightly's back_box(), requires feature asm
    // unsafe { asm!("" : : "r"(&DEFERED_MARKER)) };

    // sets panic hook to abort
    std::panic::set_hook(Box::new(|_| {
        std::process::abort();
    }));

    let mut input = vec![];

    // initialize forkserver there
    unsafe { __afl_manual_init() };

    while unsafe { __afl_persistent_loop(1000) } != 0 {
        // get buffer from AFL through stdin
        let result = io::stdin().read_to_end(&mut input);
        if result.is_err() {
            return;
        }

        // We still catch unwinding panics just in case the fuzzed code modifies
        // the panic hook.
        // If so, the fuzzer will be unable to tell different bugs appart and you will
        // only be able to find one bug at a time before fixing it to then find a new one.
        let did_panic = std::panic::catch_unwind(|| {
            closure(&input);
        })
        .is_err();

        if did_panic {
            // hopefully the custom panic hook will be called before and abort the
            // process before the stack frames are unwinded.
            std::process::abort();
        }
        input.clear();
    }
}

/// Fuzz a closure-like block of code by passing it an object of arbitrary type.
///
/// You can choose the type of the argument using the syntax as in the example below.
/// Please check out the `arbitrary` crate to see which types are available.
///
/// For performance reasons, it is recommended that you use the native type `&[u8]` when possible.
///
/// ```rust,no_run
/// # #[macro_use] extern crate afl;
/// # fn main() {
/// fuzz!(|data: &[u8]| {
///     if data.len() != 6 {return}
///     if data[0] != b'q' {return}
///     if data[1] != b'w' {return}
///     if data[2] != b'e' {return}
///     if data[3] != b'r' {return}
///     if data[4] != b't' {return}
///     if data[5] != b'y' {return}
///     panic!("BOOM")
/// });
/// # }
/// ```
#[macro_export]
macro_rules! fuzz {
    (|$buf:ident| $body:block) => {
        afl::fuzz(|$buf| $body);
    };
    (|$buf:ident: &[u8]| $body:block) => {
        afl::fuzz(|$buf| $body);
    };
    (|$buf:ident: $dty: ty| $body:block) => {
        afl::fuzz(|$buf| {
            let $buf: $dty = {
                use arbitrary::{Arbitrary, RingBuffer};
                if let Ok(d) = RingBuffer::new($buf, $buf.len())
                    .and_then(|mut b| Arbitrary::arbitrary(&mut b).map_err(|_| ""))
                {
                    d
                } else {
                    return;
                }
            };

            $body
        });
    };
}

#[cfg(test)]
mod test {
    /*
    use std::path::PathBuf;
    use std::process::{Command, Stdio};
    use std::thread;
    use std::time;

    extern crate libc;
    extern crate tempdir;

    fn target_path() -> PathBuf {
        if PathBuf::from("../target/debug/cargo-afl-fuzz").exists() {
            PathBuf::from("../target/debug/")
        } else if PathBuf::from("target/debug/cargo-afl-fuzz").exists() {
            PathBuf::from("target/debug/")
        } else {
            panic!("Could not find cargo-afl-fuzz!");
        }
    }

    #[test]
    fn test_cargo_afl_fuzz() {
        let temp_dir = tempdir::TempDir::new("aflrs").expect("Could not create temporary directory");
        let temp_dir_path = temp_dir.path();
        let mut child = Command::new(target_path().join("cargo-afl-fuzz"))
            .stdout(Stdio::inherit())
            .stderr(Stdio::inherit())
            .arg("-i")
            .arg(".")
            .arg("-o")
            .arg(temp_dir_path)
            .arg(target_path().join("examples").join("hello"))
            .spawn()
            .expect("Could not run cargo-afl-fuzz");
        thread::sleep(time::Duration::from_secs(7));
        child.kill().unwrap();
        assert!(temp_dir_path.join("fuzzer_stats").is_file());
    }
    */
}
<commit_msg>Add a fuzz! version that doesn't hook panics (#154)<commit_after>// Copyright 2015 Keegan McAllister.
// Copyright 2016 Corey Farwell.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// See `LICENSE` in this repository.

use std::io::{self, Read};
use std::{panic, process};

/// Utility that reads a `Vec` of bytes from standard input (stdin)
/// and passes it to `closure`. All panics that occur within
/// `closure` will be treated as aborts. This is done so that
/// AFL considers a panic to be a crash.
///
/// # Examples
///
/// ```no_run
/// extern crate afl;
/// # struct Image;
/// # impl Image {
/// #     fn parse(_: &[u8]) {}
/// # }
///
/// fn main() {
///     afl::read_stdio_bytes(|read| {
///         Image::parse(&read)
///     })
/// }
/// ```
#[deprecated(
    since = "0.3.3",
    note = "This function does not use the `persistent mode` and `defered forkserver mode` and is therefore very slow.
Please use fuzz() or fuzz!() instead."
)]
pub fn read_stdio_bytes<F>(closure: F)
where
    F: Fn(Vec<u8>) + panic::RefUnwindSafe,
{
    let mut input = vec![];
    let result = io::stdin().read_to_end(&mut input);
    if result.is_err() {
        return;
    }
    let was_panic = panic::catch_unwind(|| {
        closure(input);
    });
    if was_panic.is_err() {
        process::abort();
    }
}

/// Utility that reads a `String` from standard input (stdin) and
/// passes it to `closure`. If a `String` cannot be constructed from
/// the data provided by standard input, `closure` will _not_ be
/// called. All panics that occur within `closure` will be treated as
/// aborts. This is done so that AFL considers a panic to be a crash.
///
/// # Examples
///
/// ```no_run
/// extern crate afl;
/// # struct Url;
/// # impl Url {
/// #     fn parse(_: &str) {}
/// # }
///
/// fn main() {
///     afl::read_stdio_string(|string| {
///         Url::parse(&string)
///     })
/// }
/// ```
#[deprecated(
    since = "0.3.3",
    note = "This function does not use the `persistent mode` and `defered forkserver mode` and is therefore very slow.
Please use fuzz() or fuzz!() instead."
)]
pub fn read_stdio_string<F>(closure: F)
where
    F: Fn(String) + panic::RefUnwindSafe,
{
    let mut input = String::new();
    let result = io::stdin().read_to_string(&mut input);
    if result.is_err() {
        return;
    }
    let was_panic = panic::catch_unwind(|| {
        closure(input);
    });
    if was_panic.is_err() {
        process::abort();
    }
}

// those functions are provided by the afl-llvm-rt static library
extern "C" {
    fn __afl_persistent_loop(counter: usize) -> isize;
    fn __afl_manual_init();
}

/// Fuzz a closure by passing it a `&[u8]`
///
/// This slice contains a "random" quantity of "random" data.
///
/// ```rust,no_run
/// # extern crate afl;
/// # use afl::fuzz;
/// # fn main() {
/// fuzz(true, |data|{
///     if data.len() != 6 {return}
///     if data[0] != b'q' {return}
///     if data[1] != b'w' {return}
///     if data[2] != b'e' {return}
///     if data[3] != b'r' {return}
///     if data[4] != b't' {return}
///     if data[5] != b'y' {return}
///     panic!("BOOM")
/// });
/// # }
/// ```
pub fn fuzz<F>(hook: bool, closure: F)
where
    F: Fn(&[u8]) + std::panic::RefUnwindSafe,
{
    // this marker strings needs to be in the produced executable for
    // afl-fuzz to detect `persistent mode` and `defered mode`
    static PERSIST_MARKER: &'static str = "##SIG_AFL_PERSISTENT##\0";
    static DEFERED_MARKER: &'static str = "##SIG_AFL_DEFER_FORKSRV##\0";

    // we now need a fake instruction to prevent the compiler from optimizing out
    // those marker strings
    unsafe { std::ptr::read_volatile(&PERSIST_MARKER) }; // hack used in https://github.com/bluss/bencher for black_box()
    unsafe { std::ptr::read_volatile(&DEFERED_MARKER) };
    // unsafe { asm!("" : : "r"(&PERSIST_MARKER)) }; // hack used in nightly's back_box(), requires feature asm
    // unsafe { asm!("" : : "r"(&DEFERED_MARKER)) };

    if hook {
      // sets panic hook to abort
      std::panic::set_hook(Box::new(|_| {
          std::process::abort();
      }));
    }

    let mut input = vec![];

    // initialize forkserver there
    unsafe { __afl_manual_init() };

    while unsafe { __afl_persistent_loop(1000) } != 0 {
        // get buffer from AFL through stdin
        let result = io::stdin().read_to_end(&mut input);
        if result.is_err() {
            return;
        }

        // We still catch unwinding panics just in case the fuzzed code modifies
        // the panic hook.
        // If so, the fuzzer will be unable to tell different bugs appart and you will
        // only be able to find one bug at a time before fixing it to then find a new one.
        let did_panic = std::panic::catch_unwind(|| {
            closure(&input);
        })
        .is_err();

        if did_panic {
            // hopefully the custom panic hook will be called before and abort the
            // process before the stack frames are unwinded.
            std::process::abort();
        }
        input.clear();
    }
}

/// Fuzz a closure-like block of code by passing it an object of arbitrary type.
///
/// You can choose the type of the argument using the syntax as in the example below.
/// Please check out the `arbitrary` crate to see which types are available.
///
/// For performance reasons, it is recommended that you use the native type `&[u8]` when possible.
///
/// ```rust,no_run
/// # #[macro_use] extern crate afl;
/// # fn main() {
/// fuzz!(|data: &[u8]| {
///     if data.len() != 6 {return}
///     if data[0] != b'q' {return}
///     if data[1] != b'w' {return}
///     if data[2] != b'e' {return}
///     if data[3] != b'r' {return}
///     if data[4] != b't' {return}
///     if data[5] != b'y' {return}
///     panic!("BOOM")
/// });
/// # }
/// ```
#[macro_export]
macro_rules! fuzz {
    (|$buf:ident| $body:block) => {
        afl::fuzz(true, |$buf| $body);
    };
    (|$buf:ident: &[u8]| $body:block) => {
        afl::fuzz(true, |$buf| $body);
    };
    (|$buf:ident: $dty: ty| $body:block) => {
        afl::fuzz(true, |$buf| {
            let $buf: $dty = {
                use arbitrary::{Arbitrary, RingBuffer};
                if let Ok(d) = RingBuffer::new($buf, $buf.len()).and_then(|mut b|{
                        Arbitrary::arbitrary(&mut b).map_err(|_| "")
                    }) {
                    d
                } else {
                    return
                }
            };

            $body
        });
    };
}

/// Fuzz a closure-like block of code by passing it an object of arbitrary type. Panics that are
/// caught inside the fuzzed code are not turned into crashes.
///
/// You can choose the type of the argument using the syntax as in the example below.
/// Please check out the `arbitrary` crate to see which types are available.
///
/// For performance reasons, it is recommended that you use the native type `&[u8]` when possible.
///
/// ```rust,no_run
/// # #[macro_use] extern crate afl;
/// # fn main() {
/// fuzz!(|data: &[u8]| {
///     if data.len() != 6 {return}
///     if data[0] != b'q' {return}
///     if data[1] != b'w' {return}
///     if data[2] != b'e' {return}
///     if data[3] != b'r' {return}
///     if data[4] != b't' {return}
///     if data[5] != b'y' {return}
///     panic!("BOOM")
/// });
/// # }
/// ```
#[macro_export]
macro_rules! fuzz_nohook {
    (|$buf:ident| $body:block) => {
        afl::fuzz(false, |$buf| $body);
    };
    (|$buf:ident: &[u8]| $body:block) => {
        afl::fuzz(false, |$buf| $body);
    };
    (|$buf:ident: $dty: ty| $body:block) => {
        afl::fuzz(false, |$buf| {
            let $buf: $dty = {
                use arbitrary::{Arbitrary, RingBuffer};
                if let Ok(d) = RingBuffer::new($buf, $buf.len())
                    .and_then(|mut b| Arbitrary::arbitrary(&mut b).map_err(|_| ""))
                {
                    d
                } else {
                    return;
                }
            };

            $body
        });
    };
}

#[cfg(test)]
mod test {
    /*
    use std::path::PathBuf;
    use std::process::{Command, Stdio};
    use std::thread;
    use std::time;

    extern crate libc;
    extern crate tempdir;

    fn target_path() -> PathBuf {
        if PathBuf::from("../target/debug/cargo-afl-fuzz").exists() {
            PathBuf::from("../target/debug/")
        } else if PathBuf::from("target/debug/cargo-afl-fuzz").exists() {
            PathBuf::from("target/debug/")
        } else {
            panic!("Could not find cargo-afl-fuzz!");
        }
    }

    #[test]
    fn test_cargo_afl_fuzz() {
        let temp_dir = tempdir::TempDir::new("aflrs").expect("Could not create temporary directory");
        let temp_dir_path = temp_dir.path();
        let mut child = Command::new(target_path().join("cargo-afl-fuzz"))
            .stdout(Stdio::inherit())
            .stderr(Stdio::inherit())
            .arg("-i")
            .arg(".")
            .arg("-o")
            .arg(temp_dir_path)
            .arg(target_path().join("examples").join("hello"))
            .spawn()
            .expect("Could not run cargo-afl-fuzz");
        thread::sleep(time::Duration::from_secs(7));
        child.kill().unwrap();
        assert!(temp_dir_path.join("fuzzer_stats").is_file());
    }
    */
}
<|endoftext|>


<commit_before>use std::marker::PhantomData;
use std::mem;
use std::ops::{Index, IndexMut};

#[derive(Debug, Copy, Clone)]
pub struct _0;
#[derive(Debug, Copy, Clone)]
pub struct _1;

/// Single type-level bit, `_0` or `_1`.
pub trait Bit { }
impl Bit for _0 { }
impl Bit for _1 { }

/// Nonnegative type-level integer, e.g., `((_1,_0),_1) = 0b101 = 25`.
pub trait Nat {
    fn reify() -> u64;
}
impl Nat for _0 { fn reify() -> u64 { 0 } }
impl Nat for _1 { fn reify() -> u64 { 1 } }
impl<N: Nat> Nat for (N, _0) {
    fn reify() -> u64 { N::reify() << 1 }
}
impl<N: Nat> Nat for (N, _1) {
    fn reify() -> u64 { (N::reify() << 1) | 1 }
}

/// Trait making GenericArray work
pub trait ArrayLength<T> : Nat {
	/// Associated type representing the array type for the number
	type ArrayType;
}

/// Empty array - needed to end recursion
#[allow(dead_code)]
pub struct EmptyArray<T> {
	_marker: PhantomData<T>
}

/// Array with a single element - for _1
#[allow(dead_code)]
pub struct UnitArray<T> {
	data: T
}

impl<T> ArrayLength<T> for _0 {
	type ArrayType = EmptyArray<T>;
}
impl<T> ArrayLength<T> for _1 {
	type ArrayType = UnitArray<T>;
}

#[allow(dead_code)]
pub struct GenericArrayImplEven<T, U> {
	parent1: U,
	parent2: U,
	_marker: PhantomData<T>
}

#[allow(dead_code)]
pub struct GenericArrayImplOdd<T, U> {
	parent1: U,
	parent2: U,
	data: T
}

impl<T, N: ArrayLength<T>> ArrayLength<T> for (N, _0) {
	type ArrayType = GenericArrayImplEven<T, N::ArrayType>;
}

impl<T, N: ArrayLength<T>> ArrayLength<T> for (N, _1) {
	type ArrayType = GenericArrayImplOdd<T, N::ArrayType>;
}

#[allow(dead_code)]
pub struct GenericArray<T, U: ArrayLength<T>> {
	data: U::ArrayType
}

impl<T, N> Index<usize> for GenericArray<T, N> where N: ArrayLength<T> {
	type Output = T;

	fn index(&self, i: usize) -> &T {
		assert!(i < N::reify() as usize);
		let p: *const T = self as *const GenericArray<T, N> as *const T;
		unsafe { &*p.offset(i as isize) }
	}
}

impl<T, N> IndexMut<usize> for GenericArray<T, N> where N: ArrayLength<T> {

	fn index_mut(&mut self, i: usize) -> &mut T {
		assert!(i < N::reify() as usize);
		let p: *mut T = self as *mut GenericArray<T, N> as *mut T;
		unsafe { &mut *p.offset(i as isize) }
	}
}

impl<T: Clone, N> GenericArray<T, N> where N: ArrayLength<T> {

	pub fn new() -> GenericArray<T, N> {
		unsafe { mem::zeroed() }
	}

	pub fn new_list(list: &[T]) -> GenericArray<T, N> {
		assert_eq!(list.len(), N::reify() as usize);
		let mut res = GenericArray::new();
		for i in 0..N::reify() as usize {
			res[i] = list[i].clone();
		}
		res
	}

}

#[cfg(test)]
mod test {
	use super::{_0, _1, GenericArray};

	type P97 = ((((((_1, _1), _0), _0), _0), _0), _1);

	#[test]
	fn test() {
		let mut list97 = [0; 97];
		for i in 0..97 {
			list97[i] = i as i32;
		}
	    let l : GenericArray<i32, P97> = GenericArray::new_list(&list97);
	    assert_eq!(l[0], 0);
	    assert_eq!(l[1], 1);
	    assert_eq!(l[32], 32);
	    assert_eq!(l[56], 56);
	}
}<commit_msg>Use #[repr(C)] to ensure that the correct layout is used.<commit_after>use std::marker::PhantomData;
use std::mem;
use std::ops::{Index, IndexMut};

#[derive(Debug, Copy, Clone)]
pub struct _0;
#[derive(Debug, Copy, Clone)]
pub struct _1;

/// Single type-level bit, `_0` or `_1`.
pub trait Bit { }
impl Bit for _0 { }
impl Bit for _1 { }

/// Nonnegative type-level integer, e.g., `((_1,_0),_1) = 0b101 = 25`.
pub trait Nat {
    fn reify() -> u64;
}
impl Nat for _0 { fn reify() -> u64 { 0 } }
impl Nat for _1 { fn reify() -> u64 { 1 } }
impl<N: Nat> Nat for (N, _0) {
    fn reify() -> u64 { N::reify() << 1 }
}
impl<N: Nat> Nat for (N, _1) {
    fn reify() -> u64 { (N::reify() << 1) | 1 }
}

/// Trait making GenericArray work
pub trait ArrayLength<T> : Nat {
	/// Associated type representing the array type for the number
	type ArrayType;
}

/// Empty array - needed to end recursion
#[allow(dead_code)]
pub struct EmptyArray<T> {
	_marker: PhantomData<T>
}

/// Array with a single element - for _1
#[allow(dead_code)]
#[repr(C)]
pub struct UnitArray<T> {
	data: T
}

impl<T> ArrayLength<T> for _0 {
	type ArrayType = EmptyArray<T>;
}
impl<T> ArrayLength<T> for _1 {
	type ArrayType = UnitArray<T>;
}

#[allow(dead_code)]
#[repr(C)]
pub struct GenericArrayImplEven<T, U> {
	parent1: U,
	parent2: U,
	_marker: PhantomData<T>
}

#[allow(dead_code)]
#[repr(C)]
pub struct GenericArrayImplOdd<T, U> {
	parent1: U,
	parent2: U,
	data: T
}

impl<T, N: ArrayLength<T>> ArrayLength<T> for (N, _0) {
	type ArrayType = GenericArrayImplEven<T, N::ArrayType>;
}

impl<T, N: ArrayLength<T>> ArrayLength<T> for (N, _1) {
	type ArrayType = GenericArrayImplOdd<T, N::ArrayType>;
}

#[allow(dead_code)]
pub struct GenericArray<T, U: ArrayLength<T>> {
	data: U::ArrayType
}

impl<T, N> Index<usize> for GenericArray<T, N> where N: ArrayLength<T> {
	type Output = T;

	fn index(&self, i: usize) -> &T {
		assert!(i < N::reify() as usize);
		let p: *const T = self as *const GenericArray<T, N> as *const T;
		unsafe { &*p.offset(i as isize) }
	}
}

impl<T, N> IndexMut<usize> for GenericArray<T, N> where N: ArrayLength<T> {

	fn index_mut(&mut self, i: usize) -> &mut T {
		assert!(i < N::reify() as usize);
		let p: *mut T = self as *mut GenericArray<T, N> as *mut T;
		unsafe { &mut *p.offset(i as isize) }
	}
}

impl<T: Clone, N> GenericArray<T, N> where N: ArrayLength<T> {

	pub fn new() -> GenericArray<T, N> {
		unsafe { mem::zeroed() }
	}

	pub fn new_list(list: &[T]) -> GenericArray<T, N> {
		assert_eq!(list.len(), N::reify() as usize);
		let mut res = GenericArray::new();
		for i in 0..N::reify() as usize {
			res[i] = list[i].clone();
		}
		res
	}

}

#[cfg(test)]
mod test {
	use super::{_0, _1, GenericArray};

	type P97 = ((((((_1, _1), _0), _0), _0), _0), _1);

	#[test]
	fn test() {
		let mut list97 = [0; 97];
		for i in 0..97 {
			list97[i] = i as i32;
		}
	    let l : GenericArray<i32, P97> = GenericArray::new_list(&list97);
	    assert_eq!(l[0], 0);
	    assert_eq!(l[1], 1);
	    assert_eq!(l[32], 32);
	    assert_eq!(l[56], 56);
	}
}
<|endoftext|>


<commit_before>#![feature(plugin_registrar, macro_rules)]

extern crate rustc;
extern crate syntax;

use std::io::Command;
use std::io::fs::PathExtensions;
use std::os;
use std::str;
use std::sync::{Once, ONCE_INIT};

use rustc::plugin::Registry;
use syntax::abi;
use syntax::ast;
use syntax::attr;
use syntax::ptr::P;
use syntax::codemap::Span;
use syntax::ext::base::{ExtCtxt, DummyResult};
use syntax::ext::base::MacResult;
use syntax::ext::build::AstBuilder;
use syntax::fold::Folder;
use syntax::parse::parser::Parser;
use syntax::parse::token;
use syntax::parse::token::{special_idents, InternedString};
use syntax::print::pprust;
use syntax::util::small_vector::SmallVector;

struct LibInfo {
    lib: String,
    state: State,
    deps: Vec<(String, bool)>, // (name, static)
    locs: Vec<Path>,
}

struct MacItems {
    items: Vec<P<ast::Item>>,
}

enum State { Dynamic, Static(SystemDeps) }
enum SystemDeps { SystemDynamic, SystemStatic }
enum Favor { FavorDynamic, FavorStatic }

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_macro("link_config", expand_link_config);
}

fn expand_link_config(ecx: &mut ExtCtxt, span: Span,
                      tts: &[ast::TokenTree]) -> Box<MacResult+'static> {
    macro_rules! try_dummy( ($e:expr) => (
        match $e { Ok(s) => s, Err(()) => return DummyResult::any(span) }
    ) )

    let mut parser = ecx.new_parser_from_tts(tts);
    let (pkg, sp) = try_dummy!(parse_string(ecx, &mut parser));
    let mut favor_dynamic = FavorDynamic;
    let mut dylib_state = Some(Dynamic);
    let mut static_state = Some(Static(SystemDynamic));
    if parser.eat(&token::Comma) && parser.eat(&token::OpenDelim(token::Bracket)) {
        while !parser.eat(&token::CloseDelim(token::Bracket)) {
            parser.eat(&token::Comma);
            let (modifier, sp) = try_dummy!(parse_string(ecx, &mut parser));
            match modifier.as_slice() {
                "only_static" => {
                    dylib_state = None;
                    favor_dynamic = FavorStatic;
                }
                "only_dylib" => {
                    static_state = None;
                    favor_dynamic = FavorDynamic;
                }
                "system_static" => static_state = Some(Static(SystemStatic)),
                "favor_static" => favor_dynamic = FavorStatic,
                s => ecx.span_err(sp, format!("unknown modifier: `{}`",
                                              s).as_slice()),
            }
        }
    }
    if !parser.eat(&token::Eof) {
        ecx.span_err(parser.span, "only one string literal allowed");
        return DummyResult::any(span);
    }

    let dyn = try_dummy!(system_pkgconfig(ecx, sp, pkg.as_slice(), dylib_state));
    let statik = try_dummy!(system_pkgconfig(ecx, sp, pkg.as_slice(), static_state));

    let mut items = Vec::new();
    match dyn {
        Some(info) => items.push(block(ecx, sp, &info, favor_dynamic)),
        None => {}
    }
    match statik {
        Some(info) => items.push(block(ecx, sp, &info, favor_dynamic)),
        None => {}
    }
    box MacItems { items: items } as Box<MacResult>
}

fn system_pkgconfig(ecx: &mut ExtCtxt, sp: Span, pkg: &str,
                    state: Option<State>) -> Result<Option<LibInfo>, ()> {
    let state = match state {
        Some(state) => state,
        None => return Ok(None),
    };
    add_cargo_pkg_config_paths();

    let mut cmd = Command::new("pkg-config");
    cmd.arg(pkg).arg("--libs").env("PKG_CONFIG_ALLOW_SYSTEM_LIBS", "1");
    match state {
        Static(..) => { cmd.arg("--static"); }
        Dynamic => {}
    }
    let out = match cmd.output() {
        Ok(out) => out,
        Err(e) => {
            ecx.span_err(sp, format!("could not run pkg-config: {}", e).as_slice());
            return Err(())
        }
    };
    let stdout = str::from_utf8(out.output.as_slice()).unwrap();
    let stderr = str::from_utf8(out.error.as_slice()).unwrap();
    if !out.status.success() {
        let mut msg = format!("pkg-config did not exit successfully: {}",
                              out.status);
        if stdout.len() > 0 {
            msg.push_str("\n--- stdout\n");
            msg.push_str(stdout);
        }
        if stderr.len() > 0 {
            msg.push_str("\n--- stderr\n");
            msg.push_str(stderr);
        }
        ecx.span_err(sp, msg.as_slice());
        return Err(())
    }

    let mut libs = Vec::new();
    let mut locs = Vec::new();
    for arg in stdout.split(' ').filter(|l| !l.is_empty()) {
        if arg.starts_with("-l") {
            libs.push(arg.slice_from(2));
        } else if arg.starts_with("-L") {
            locs.push(Path::new(arg.slice_from(2).to_string()));
        }
    }

    let allow_static = match state {
        Static(..) => true,
        _ => false,
    };
    let allow_static_system = match state {
        Static(SystemStatic) => true,
        _ => false,
    };

    let cargo_locs = cargo_native_dirs();
    let libs = libs.into_iter().map(|lib| {
        let mut candidates = cargo_locs.iter().chain(locs.iter());
        (lib.to_string(), allow_static && candidates.any(|base| {
            (allow_static_system || !base.as_vec().starts_with(b"/usr")) &&
                (base.join(format!("lib{}.a", lib)).exists() ||
                 base.join(format!("{}.lib", lib)).exists() ||
                 base.join(format!("lib{}.lib", lib)).exists())
        }))
    }).collect();
    Ok(Some(LibInfo {
        lib: pkg.to_string(),
        deps: libs,
        locs: locs,
        state: state,
    }))
}

fn block(ecx: &mut ExtCtxt, sp: Span, info: &LibInfo,
         favor: Favor) -> P<ast::Item> {
    let lib = token::intern_and_get_ident(info.lib.as_slice());
    let s = match favor {
        FavorDynamic => InternedString::new("statik"),
        FavorStatic => InternedString::new("dylib"),
    };
    let cfg = ecx.meta_name_value(sp, s, ast::LitStr(lib, ast::CookedStr));
    let cfg = match (info.state, favor) {
        (Static(..), FavorDynamic) |
        (Dynamic, FavorStatic) => cfg,
        (Dynamic, FavorDynamic) |
        (Static(..), FavorStatic) => {
            ecx.meta_list(sp, InternedString::new("not"), vec![cfg])
        }
    };
    let cfg = ecx.meta_list(sp, InternedString::new("cfg"), vec![cfg]);

    let attrs = vec![ecx.attribute(sp, cfg)].into_iter();
    let attrs = attrs.chain(info.deps.iter().map(|&(ref l, statik)| {
        // Build #[link(name = <l>)]
        let l = token::intern_and_get_ident(l.as_slice());
        let mut items = Vec::new();
        items.push(ecx.meta_name_value(sp, InternedString::new("name"),
                                       ast::LitStr(l, ast::CookedStr)));
        if statik {
            let l = InternedString::new("static");
            items.push(ecx.meta_name_value(sp, InternedString::new("kind"),
                                           ast::LitStr(l, ast::CookedStr)));
        }
        let list = ecx.meta_list(sp, InternedString::new("link"), items);
        ecx.attribute(sp, list)
    }));
    let attrs = attrs.chain(info.locs.iter().map(|l| {
        let l = token::intern_and_get_ident(format!("-L{}", l.display()).as_slice());
        let args = ecx.meta_name_value(sp, InternedString::new("link_args"),
                                       ast::LitStr(l, ast::CookedStr));
        ecx.attribute(sp, args)
    }));
    let mut attrs = attrs.map(|attr| {
        attr::mark_used(&attr);
        attr
    });

    ecx.item(sp, special_idents::invalid, attrs.collect(),
             ast::ItemForeignMod(ast::ForeignMod {
        abi: abi::C,
        view_items: Vec::new(),
        items: Vec::new(),
    }))
}

fn parse_string(ecx: &mut ExtCtxt,
                parser: &mut Parser) -> Result<(String, Span), ()> {
    let entry = ecx.expander().fold_expr(parser.parse_expr());
    match entry.node {
        ast::ExprLit(ref lit) => {
            match lit.node {
                ast::LitStr(ref s, _) => return Ok((s.to_string(), entry.span)),
                _ => {}
            }
        }
        _ => {}
    }
    ecx.span_err(entry.span, format!(
        "expected string literal but got `{}`",
        pprust::expr_to_string(&*entry)).as_slice());
    Err(())
}

impl MacResult for MacItems {
    fn make_items(self: Box<MacItems>) -> Option<SmallVector<P<ast::Item>>> {
        Some(self.items.into_iter().collect())
    }
}

// lol hax
fn cargo_native_dirs() -> Vec<Path> {
    match rustc::driver::handle_options(os::args()) {
        Some(matches) => {
            matches.opt_strs("L").into_iter().filter_map(|s| {
                if s.as_slice().contains("native") {
                    Some(Path::new(s))
                } else {
                    None
                }
            }).collect()
        }
        None => Vec::new()
    }
}

fn add_cargo_pkg_config_paths() {
    static mut DONE: Once = ONCE_INIT;
    unsafe { DONE.doit(add_cargo_pkg_config_paths) }

    fn add_cargo_pkg_config_paths() {
        let path = os::getenv_as_bytes("PKG_CONFIG_PATH").unwrap_or(Vec::new());
        let mut pkg_config_path = os::split_paths(path.as_slice());
        pkg_config_path.push_all(cargo_native_dirs().as_slice());
        let pkg_config_path = os::join_paths(pkg_config_path.as_slice())
                                 .unwrap();
        os::setenv("PKG_CONFIG_PATH", pkg_config_path.as_slice());
    }
}
<commit_msg>Fixes enum syntax and handle_options<commit_after>#![feature(plugin_registrar, macro_rules)]

extern crate rustc;
extern crate rustc_trans;
extern crate syntax;

use std::io::Command;
use std::io::fs::PathExtensions;
use std::os;
use std::str;
use std::sync::{Once, ONCE_INIT};

use rustc::plugin::Registry;
use syntax::abi;
use syntax::ast;
use syntax::attr;
use syntax::ptr::P;
use syntax::codemap::Span;
use syntax::ext::base::{ExtCtxt, DummyResult};
use syntax::ext::base::MacResult;
use syntax::ext::build::AstBuilder;
use syntax::fold::Folder;
use syntax::parse::parser::Parser;
use syntax::parse::token;
use syntax::parse::token::{special_idents, InternedString};
use syntax::print::pprust;
use syntax::util::small_vector::SmallVector;

struct LibInfo {
    lib: String,
    state: State,
    deps: Vec<(String, bool)>, // (name, static)
    locs: Vec<Path>,
}

struct MacItems {
    items: Vec<P<ast::Item>>,
}

enum State { Dynamic, Static(SystemDeps) }
enum SystemDeps { SystemDynamic, SystemStatic }
enum Favor { FavorDynamic, FavorStatic }

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_macro("link_config", expand_link_config);
}

fn expand_link_config(ecx: &mut ExtCtxt, span: Span,
                      tts: &[ast::TokenTree]) -> Box<MacResult+'static> {
    macro_rules! try_dummy( ($e:expr) => (
        match $e { Ok(s) => s, Err(()) => return DummyResult::any(span) }
    ) )

    let mut parser = ecx.new_parser_from_tts(tts);
    let (pkg, sp) = try_dummy!(parse_string(ecx, &mut parser));
    let mut favor_dynamic = Favor::FavorDynamic;
    let mut dylib_state = Some(State::Dynamic);
    let mut static_state = Some(State::Static(SystemDeps::SystemDynamic));
    if parser.eat(&token::Comma) && parser.eat(&token::OpenDelim(token::Bracket)) {
        while !parser.eat(&token::CloseDelim(token::Bracket)) {
            parser.eat(&token::Comma);
            let (modifier, sp) = try_dummy!(parse_string(ecx, &mut parser));
            match modifier.as_slice() {
                "only_static" => {
                    dylib_state = None;
                    favor_dynamic = Favor::FavorStatic;
                }
                "only_dylib" => {
                    static_state = None;
                    favor_dynamic = Favor::FavorDynamic;
                }
                "system_static" => static_state = Some(State::Static(SystemDeps::SystemStatic)),
                "favor_static" => favor_dynamic = Favor::FavorStatic,
                s => ecx.span_err(sp, format!("unknown modifier: `{}`",
                                              s).as_slice()),
            }
        }
    }
    if !parser.eat(&token::Eof) {
        ecx.span_err(parser.span, "only one string literal allowed");
        return DummyResult::any(span);
    }

    let dyn = try_dummy!(system_pkgconfig(ecx, sp, pkg.as_slice(), dylib_state));
    let statik = try_dummy!(system_pkgconfig(ecx, sp, pkg.as_slice(), static_state));

    let mut items = Vec::new();
    match dyn {
        Some(info) => items.push(block(ecx, sp, &info, favor_dynamic)),
        None => {}
    }
    match statik {
        Some(info) => items.push(block(ecx, sp, &info, favor_dynamic)),
        None => {}
    }
    box MacItems { items: items } as Box<MacResult>
}

fn system_pkgconfig(ecx: &mut ExtCtxt, sp: Span, pkg: &str,
                    state: Option<State>) -> Result<Option<LibInfo>, ()> {
    let state = match state {
        Some(state) => state,
        None => return Ok(None),
    };
    add_cargo_pkg_config_paths();

    let mut cmd = Command::new("pkg-config");
    cmd.arg(pkg).arg("--libs").env("PKG_CONFIG_ALLOW_SYSTEM_LIBS", "1");
    match state {
        State::Static(..) => { cmd.arg("--static"); }
        State::Dynamic => {}
    }
    let out = match cmd.output() {
        Ok(out) => out,
        Err(e) => {
            ecx.span_err(sp, format!("could not run pkg-config: {}", e).as_slice());
            return Err(())
        }
    };
    let stdout = str::from_utf8(out.output.as_slice()).unwrap();
    let stderr = str::from_utf8(out.error.as_slice()).unwrap();
    if !out.status.success() {
        let mut msg = format!("pkg-config did not exit successfully: {}",
                              out.status);
        if stdout.len() > 0 {
            msg.push_str("\n--- stdout\n");
            msg.push_str(stdout);
        }
        if stderr.len() > 0 {
            msg.push_str("\n--- stderr\n");
            msg.push_str(stderr);
        }
        ecx.span_err(sp, msg.as_slice());
        return Err(())
    }

    let mut libs = Vec::new();
    let mut locs = Vec::new();
    for arg in stdout.split(' ').filter(|l| !l.is_empty()) {
        if arg.starts_with("-l") {
            libs.push(arg.slice_from(2));
        } else if arg.starts_with("-L") {
            locs.push(Path::new(arg.slice_from(2).to_string()));
        }
    }

    let allow_static = match state {
        State::Static(..) => true,
        _ => false,
    };
    let allow_static_system = match state {
        State::Static(SystemStatic) => true,
        _ => false,
    };

    let cargo_locs = cargo_native_dirs();
    let libs = libs.into_iter().map(|lib| {
        let mut candidates = cargo_locs.iter().chain(locs.iter());
        (lib.to_string(), allow_static && candidates.any(|base| {
            (allow_static_system || !base.as_vec().starts_with(b"/usr")) &&
                (base.join(format!("lib{}.a", lib)).exists() ||
                 base.join(format!("{}.lib", lib)).exists() ||
                 base.join(format!("lib{}.lib", lib)).exists())
        }))
    }).collect();
    Ok(Some(LibInfo {
        lib: pkg.to_string(),
        deps: libs,
        locs: locs,
        state: state,
    }))
}

fn block(ecx: &mut ExtCtxt, sp: Span, info: &LibInfo,
         favor: Favor) -> P<ast::Item> {
    let lib = token::intern_and_get_ident(info.lib.as_slice());
    let s = match favor {
        Favor::FavorDynamic => InternedString::new("statik"),
        Favor::FavorStatic => InternedString::new("dylib"),
    };
    let cfg = ecx.meta_name_value(sp, s, ast::LitStr(lib, ast::CookedStr));
    let cfg = match (info.state, favor) {
        (State::Static(..), Favor::FavorDynamic) |
        (State::Dynamic, Favor::FavorStatic) => cfg,
        (State::Dynamic, Favor::FavorDynamic) |
        (State::Static(..), Favor::FavorStatic) => {
            ecx.meta_list(sp, InternedString::new("not"), vec![cfg])
        }
    };
    let cfg = ecx.meta_list(sp, InternedString::new("cfg"), vec![cfg]);

    let attrs = vec![ecx.attribute(sp, cfg)].into_iter();
    let attrs = attrs.chain(info.deps.iter().map(|&(ref l, statik)| {
        // Build #[link(name = <l>)]
        let l = token::intern_and_get_ident(l.as_slice());
        let mut items = Vec::new();
        items.push(ecx.meta_name_value(sp, InternedString::new("name"),
                                       ast::LitStr(l, ast::CookedStr)));
        if statik {
            let l = InternedString::new("static");
            items.push(ecx.meta_name_value(sp, InternedString::new("kind"),
                                           ast::LitStr(l, ast::CookedStr)));
        }
        let list = ecx.meta_list(sp, InternedString::new("link"), items);
        ecx.attribute(sp, list)
    }));
    let attrs = attrs.chain(info.locs.iter().map(|l| {
        let l = token::intern_and_get_ident(format!("-L{}", l.display()).as_slice());
        let args = ecx.meta_name_value(sp, InternedString::new("link_args"),
                                       ast::LitStr(l, ast::CookedStr));
        ecx.attribute(sp, args)
    }));
    let mut attrs = attrs.map(|attr| {
        attr::mark_used(&attr);
        attr
    });

    ecx.item(sp, special_idents::invalid, attrs.collect(),
             ast::ItemForeignMod(ast::ForeignMod {
        abi: abi::C,
        view_items: Vec::new(),
        items: Vec::new(),
    }))
}

fn parse_string(ecx: &mut ExtCtxt,
                parser: &mut Parser) -> Result<(String, Span), ()> {
    let entry = ecx.expander().fold_expr(parser.parse_expr());
    match entry.node {
        ast::ExprLit(ref lit) => {
            match lit.node {
                ast::LitStr(ref s, _) => return Ok((s.to_string(), entry.span)),
                _ => {}
            }
        }
        _ => {}
    }
    ecx.span_err(entry.span, format!(
        "expected string literal but got `{}`",
        pprust::expr_to_string(&*entry)).as_slice());
    Err(())
}

impl MacResult for MacItems {
    fn make_items(self: Box<MacItems>) -> Option<SmallVector<P<ast::Item>>> {
        Some(self.items.into_iter().collect())
    }
}

// lol hax
fn cargo_native_dirs() -> Vec<Path> {
    match rustc_trans::driver::handle_options(os::args()) {
        Some(matches) => {
            matches.opt_strs("L").into_iter().filter_map(|s| {
                if s.as_slice().contains("native") {
                    Some(Path::new(s))
                } else {
                    None
                }
            }).collect()
        }
        None => Vec::new()
    }
}

fn add_cargo_pkg_config_paths() {
    static mut DONE: Once = ONCE_INIT;
    unsafe { DONE.doit(add_cargo_pkg_config_paths) }

    fn add_cargo_pkg_config_paths() {
        let path = os::getenv_as_bytes("PKG_CONFIG_PATH").unwrap_or(Vec::new());
        let mut pkg_config_path = os::split_paths(path.as_slice());
        pkg_config_path.push_all(cargo_native_dirs().as_slice());
        let pkg_config_path = os::join_paths(pkg_config_path.as_slice())
                                 .unwrap();
        os::setenv("PKG_CONFIG_PATH", pkg_config_path.as_slice());
    }
}
<|endoftext|>


<commit_before>#![feature(collections, unboxed_closures, core)]
extern crate itertools;
extern crate rustc_serialize;

use std::collections::{BTreeMap, HashMap, HashSet};
use std::collections::hash_map::Entry::*;
use std::hash::{Hash, Hasher};
use std::sync::Arc;
use std::iter;
use std::ops;

use itertools::{GroupBy, Itertools};

/// A Document contains an id and content.
/// Hashing and equality are based only on the id field.
#[derive(Clone, Eq, Debug, RustcEncodable, RustcDecodable)]
pub struct Document {
    id: String,
    content: String,
}

impl Document {
    /// Construct a new Document from an id and content.
    /// Both two arguments can be anything that can be turned into a String.
    pub fn new<S, T>(id: S, content: T) -> Document
        where S: Into<String>,
              T: Into<String>
    {
        Document { id: id.into(), content: content.into(), }
    }

    pub fn id(&self) -> &str {
        &self.id
    }

    pub fn content(&self) -> &str {
        &self.content
    }
}

impl Hash for Document {
    // Documents are unique only upon their id
    fn hash<H>(&self, state: &mut H) where H: Hasher {
        self.id.hash(state);
    }
}

impl PartialEq<Document> for Document {
    // Documents are unique only upon their id
    fn eq(&self, other: &Document) -> bool {
        self.id == other.id
    }
}

#[derive(Clone, Eq, PartialEq, Hash, Debug)]
struct IndexedDocument {
    pub doc: Arc<Document>,
    pub highlighted: (usize, usize),
}

/// A SearchResult is the representation of a Document returned for a specific set of search
/// terms. It is unique upon the document and the vec of highlight indices. It also contains a
/// score and a String of the document's content post-highlighting.
#[derive(Clone, Debug, RustcEncodable)]
pub struct SearchResult {
    doc: Arc<Document>,
    highlights: Vec<(usize, usize)>,
    highlighted: String,
    score: f32,
}

impl PartialEq for SearchResult {
    // SearchResult is unique upon its document and highlights
    fn eq(&self, other: &SearchResult) -> bool {
        self.doc == other.doc && self.highlights == other.highlights
    }
}

impl Eq for SearchResult {}

impl Hash for SearchResult {
    // SearchResult is unique upon its document and highlights
    fn hash<H>(&self, state: &mut H) where H: Hasher {
        self.doc.hash(state);
        self.highlights.hash(state);
    }
}

impl SearchResult {
    fn new(doc: Arc<Document>, mut highlights: Vec<(usize, usize)>) -> SearchResult {
        SearchResult {
            score: highlights.iter()
                .map(|&(begin, end)| end - begin)
                .sum::<usize>() as f32 / (doc.content.len() as f32).sqrt(),
            highlighted: SearchResult::highlighted_content(&doc.content, &mut highlights),
            doc: doc,
            highlights: highlights,
        }
    }

    /// Returns the document
    pub fn doc(&self) -> &Arc<Document> {
        &self.doc
    }

    /// Returns the highlighted indices.
    ///
    /// Each `(usize, usize)` indicates the start and end of a term in the document's content
    /// that should be highlighted.
    pub fn highlights(&self) -> &Vec<(usize, usize)> {
        &self.highlights
    }

    /// Returns the highlighted content.
    ///
    /// The highlighted content is the result of splicing in `<span class=highlight></span>` tags
    /// around words that should be highlighted.
    pub fn highlighted(&self) -> &String {
        &self.highlighted
    }

    /// Returns the search result's score.
    ///
    /// Score is computed by the product of the summed length of the matching terms and the inverse
    /// square root of the length of the document. Taking the square root of the document's length
    /// helps to combat bias toward short content.
    pub fn score(&self) -> f32 {
        self.score
    }

    fn highlighted_content(content: &str, highlights: &mut [(usize, usize)]) -> String {
        highlights.sort();
        let mut begin_idx = 0;
        let mut parts = vec![];
        for &mut (begin, end) in highlights {
            parts.push(&content[begin_idx..begin]);
            parts.push("");
            parts.push(&content[begin..end]);
            parts.push("");
            begin_idx = end;
        }
        parts.push(&content[begin_idx..]);
        parts.into_iter().join("")
    }
}

/// A basic implementation of an `Index`, the inverted index is a data structure that maps
/// from words to sets of Documents.
#[derive(Debug)]
pub struct InvertedIndex {
    index: BTreeMap<String, HashSet<IndexedDocument>>,
    docs: BTreeMap<String, Arc<Document>>,
}

impl InvertedIndex {
    pub fn new() -> InvertedIndex {
        InvertedIndex {
            index: BTreeMap::new(),
            docs: BTreeMap::new(),
        }
    }

    /// A basic implementation of index, splits the document's content into whitespace-separated
    /// words, and inserts each word-document pair into the map.
    pub fn index(&mut self, doc: Document) {
        let doc = Arc::new(doc);
        let analyzed  = analyze_doc(doc.content());
        let previous_version = self.docs.insert(doc.id.clone(), doc.clone());
        if let Some(previous_version) = previous_version {
            let previous_analyzed = analyze_doc(previous_version.content());
            for (ngram, highlights) in previous_analyzed {
                let is_empty = {
                    let docs_for_ngram = self.index.get_mut(&ngram).unwrap();
                    docs_for_ngram.remove(&IndexedDocument { doc: doc.clone(), highlighted: highlights });
                    docs_for_ngram.is_empty()
                };
                if is_empty {
                    self.index.remove(&ngram);
                }
            }
        }

        for (ngram, highlighted) in analyzed {
            self.index.entry(ngram).or_insert_with(|| HashSet::new())
                .insert(IndexedDocument { doc: doc.clone(), highlighted: highlighted });
        }
    }

    /// A basic search implementation that splits the query's content into whitespace-separated
    /// words, looks up the set of Documents for each word, and then concatenates the sets.
    pub fn search(&self, query: &str) -> Vec<SearchResult> {
        let unique_terms: HashSet<_> = query.split_whitespace().collect();
        let map = unique_terms.into_iter()
            .flat_map(|word| self.index.get(&word.to_lowercase()))
            .flat_map(|doc| doc)
            .cloned()
            .fold(HashMap::new(), |mut map, IndexedDocument { doc, highlighted: (begin, end) }| {
                match map.entry(doc) {
                    Vacant(entry) => {
                        let mut indices = HashMap::new();
                        indices.insert(begin, end);
                        entry.insert(indices);
                    }
                    Occupied(mut entry) => {
                        match entry.get_mut().entry(begin) {
                            Vacant(entry) => { entry.insert(end); }
                            Occupied(mut entry) => {
                                let current = entry.get_mut();
                                *current = std::cmp::max(*current, end);
                            }
                        }
                    }
                }
                map
            });
        let mut results: Vec<_> = map.into_iter()
            .map(|(doc, index_map)| SearchResult::new(doc, index_map.into_iter().collect()))
            .collect();
        results.sort_by(|result1, result2| result2.score.partial_cmp(&result1.score).unwrap());
        results
    }
}

fn analyze_doc(doc: &str)
-> iter::FlatMap<
    iter::Filter<
        GroupBy<bool, std::str::CharIndices, fn(&(usize, char)) -> bool>,
        fn(&(bool, Vec<(usize, char)>)) -> bool>,
    iter::Map<ops::Range<usize>, Ngrams>,
    fn((bool, Vec<(usize, char)>)) -> iter::Map<ops::Range<usize>, Ngrams>>
{
    doc.char_indices()
        .group_by(is_whitespace as fn(&(usize, char)) -> bool)
        .filter(not_whitespace as fn(&(bool, Vec<(usize, char)>)) -> bool)
        .flat_map(ngrams)
}

fn ngrams((_, chars): (bool, Vec<(usize, char)>)) -> iter::Map<ops::Range<usize>, Ngrams> {
    (1..chars.len() + 1).map(Ngrams::new(chars))
}

fn not_whitespace(&(is_whitespace, _): &(bool, Vec<(usize, char)>)) -> bool {
    !is_whitespace
}

fn is_whitespace(&(_, c): &(usize, char)) -> bool {
    c.is_whitespace()
}

struct Ngrams {
    chars: Vec<(usize, char)>
}

impl Ngrams {
    fn new(chars: Vec<(usize, char)>) -> Ngrams {
        Ngrams { chars: chars }
    }
}

impl Fn<(usize,)> for Ngrams {
    extern "rust-call" fn call(&self, (to,): (usize,)) -> (String, (usize, usize)) {
        let word = self.chars[..to].iter().flat_map(|&(_, c)| c.to_lowercase()).collect();
        let start = self.chars[0].0;
        let (last_idx, last_char) = self.chars[to - 1];
        let finish = last_idx + last_char.len_utf8();
        (word, (start, finish))
    }
}

impl FnMut<(usize,)> for Ngrams {
    extern "rust-call" fn call_mut(&mut self, to: (usize,)) -> (String, (usize, usize)) {
        self.call(to)
    }
}

impl FnOnce<(usize,)> for Ngrams {
    type Output = (String, (usize, usize));
    extern "rust-call" fn call_once(self, to: (usize,)) -> (String, (usize, usize)) {
        self.call(to)
    }
}

#[test]
fn test_search1() {
    let mut index = InvertedIndex::new();
    let doc1 = Document::new("1", "learn to program in rust today");
    index.index(doc1.clone());
    let doc2 = Document::new("2", "what did you today do");
    index.index(doc2.clone());
    let search_results = index.search("to");
    let expected = [
        SearchResult::new(Arc::new(doc1), vec![(6, 8), (25, 27)]),
        SearchResult::new(Arc::new(doc2), vec![(13, 15)])
    ];
    assert_eq!(search_results, expected.iter().cloned().collect::<Vec<_>>());
    assert_eq!("learn to program in rust today", expected[0].highlighted);
}

#[test]
fn test_ngrams() {
    let mut index = InvertedIndex::new();
    let doc1 = Document::new("1", "learn to program in rust today");
    let doc2 = Document::new("2", "what did you today do");
    index.index(doc1.clone());
    index.index(doc2.clone());
    let search_results = index.search("to");
    let expected = [
        SearchResult::new(Arc::new(doc1), vec![(6, 8), (25, 27)]),
        SearchResult::new(Arc::new(doc2), vec![(13, 15)]),
    ];
    assert_eq!(search_results, expected.iter().cloned().collect::<Vec<_>>());
    assert_eq!("learn to program in rust today", expected[0].highlighted);

}

#[test]
fn test_search2() {
    let mut index = InvertedIndex::new();
    let doc1 = Document::new("2", "what to do today");
    let doc2 = Document::new("3", "hey today");
    index.index(doc1.clone());
    index.index(doc2.clone());
    let search_results = index.search("to");
    let expected = [
        SearchResult::new(Arc::new(doc1), vec![(5, 7), (11, 13)]),
        SearchResult::new(Arc::new(doc2), vec![(4, 6)]),
    ];
    assert_eq!(search_results, expected.iter().cloned().collect::<Vec<_>>());
    assert_eq!(expected[0].highlighted, "what to do today");
}

#[test]
fn test_unicode() {
    let mut index = InvertedIndex::new();
    let doc = Document::new("0", "abc åäö");
    index.index(doc.clone());
    let to_search = "å";
    let search_results = index.search(to_search);
    let &SearchResult { ref doc, ref highlights, .. } = search_results.iter().next().unwrap();
    let (begin, end) = highlights[0];
    assert_eq!(&doc.content()[begin..end], to_search);
}

#[test]
fn test_update_doc() {
    let mut index = InvertedIndex::new();
    let doc = Document::new("0", "abc åäö");
    index.index(doc);
    let doc = Document::new("0", "different");
    index.index(doc);
    let search_results = index.search("å");
    assert!(search_results.is_empty());
    assert_eq!(index.docs.len(), 1);
}

#[test]
fn test_ranking() {
    let mut index = InvertedIndex::new();
    let doc = Document::new("0", "beat");
    index.index(doc.clone());
    let doc2 = Document::new("1", "beast");
    index.index(doc2.clone());
    let search_results = index.search("be");
    assert_eq!(index.docs.len(), 2);
    assert_eq!(&*search_results.into_iter().next().unwrap().doc, &doc);
}

#[test]
fn test_duplicate_term() {
    let mut index = InvertedIndex::new();
    let doc = Document::new("0", "beat");
    index.index(doc.clone());
    let search_results = index.search("be be");
    assert_eq!(search_results.len(), 1);
}

#[test]
fn test_duplicate_term2() {
    let mut index = InvertedIndex::new();
    let doc = Document::new("0", "beat");
    index.index(doc.clone());
    let search_results = index.search("be b");
    assert_eq!(search_results.len(), 1);
    assert_eq!(search_results[0], SearchResult::new(Arc::new(doc), vec![(0, 2)]));
}
<commit_msg>Remove highlight field from SearchResult struct. It probably shouldn't bake in a particular way of doing things.<commit_after>#![feature(collections, unboxed_closures, core)]
extern crate itertools;
extern crate rustc_serialize;

use std::collections::{BTreeMap, HashMap, HashSet};
use std::collections::hash_map::Entry::*;
use std::hash::{Hash, Hasher};
use std::sync::Arc;
use std::iter;
use std::ops;

use itertools::{GroupBy, Itertools};

/// A Document contains an id and content.
/// Hashing and equality are based only on the id field.
#[derive(Clone, Eq, Debug, RustcEncodable, RustcDecodable)]
pub struct Document {
    id: String,
    content: String,
}

impl Document {
    /// Construct a new Document from an id and content.
    /// Both two arguments can be anything that can be turned into a String.
    pub fn new<S, T>(id: S, content: T) -> Document
        where S: Into<String>,
              T: Into<String>
    {
        Document { id: id.into(), content: content.into(), }
    }

    pub fn id(&self) -> &str {
        &self.id
    }

    pub fn content(&self) -> &str {
        &self.content
    }
}

impl Hash for Document {
    // Documents are unique only upon their id
    fn hash<H>(&self, state: &mut H) where H: Hasher {
        self.id.hash(state);
    }
}

impl PartialEq<Document> for Document {
    // Documents are unique only upon their id
    fn eq(&self, other: &Document) -> bool {
        self.id == other.id
    }
}

#[derive(Clone, Eq, PartialEq, Hash, Debug)]
struct IndexedDocument {
    pub doc: Arc<Document>,
    pub highlighted: (usize, usize),
}

/// A SearchResult is the representation of a Document returned for a specific set of search
/// terms. It is unique upon the document and the vec of highlight indices. It also contains a
/// score and a String of the document's content post-highlighting.
#[derive(Clone, Debug, RustcEncodable)]
pub struct SearchResult {
    doc: Arc<Document>,
    highlights: Vec<(usize, usize)>,
    score: f32,
}

impl PartialEq for SearchResult {
    // SearchResult is unique upon its document and highlights
    fn eq(&self, other: &SearchResult) -> bool {
        self.doc == other.doc && self.highlights == other.highlights
    }
}

impl Eq for SearchResult {}

impl Hash for SearchResult {
    // SearchResult is unique upon its document and highlights
    fn hash<H>(&self, state: &mut H) where H: Hasher {
        self.doc.hash(state);
        self.highlights.hash(state);
    }
}

impl SearchResult {
    fn new(doc: Arc<Document>, mut highlights: Vec<(usize, usize)>) -> SearchResult {
        SearchResult {
            score: highlights.iter()
                .map(|&(begin, end)| end - begin)
                .sum::<usize>() as f32 / (doc.content.len() as f32).sqrt(),
            doc: doc,
            highlights: { highlights.sort(); highlights },
        }
    }

    /// Returns the document
    pub fn doc(&self) -> &Arc<Document> {
        &self.doc
    }

    /// Returns the highlighted indices.
    ///
    /// Each `(usize, usize)` indicates the start and end of a term in the document's content
    /// that should be highlighted.
    pub fn highlights(&self) -> &Vec<(usize, usize)> {
        &self.highlights
    }

    /// Returns the search result's score.
    ///
    /// Score is computed by the product of the summed length of the matching terms and the inverse
    /// square root of the length of the document. Taking the square root of the document's length
    /// helps to combat bias toward short content.
    pub fn score(&self) -> f32 {
        self.score
    }

    pub fn highlight(&self, before: &str, after: &str) -> String {
        let mut begin_idx = 0;
        let mut parts = String::new();
        for &(begin, end) in &self.highlights {
            parts.push_str(&self.doc.content[begin_idx..begin]);
            parts.push_str(before);
            parts.push_str(&self.doc.content[begin..end]);
            parts.push_str(after);
            begin_idx = end;
        }
        parts.push_str(&self.doc.content[begin_idx..]);
        parts
    }
}

/// A basic implementation of an `Index`, the inverted index is a data structure that maps
/// from words to sets of Documents.
#[derive(Debug)]
pub struct InvertedIndex {
    index: BTreeMap<String, HashSet<IndexedDocument>>,
    docs: BTreeMap<String, Arc<Document>>,
}

impl InvertedIndex {
    pub fn new() -> InvertedIndex {
        InvertedIndex {
            index: BTreeMap::new(),
            docs: BTreeMap::new(),
        }
    }

    /// A basic implementation of index, splits the document's content into whitespace-separated
    /// words, and inserts each word-document pair into the map.
    pub fn index(&mut self, doc: Document) {
        let doc = Arc::new(doc);
        let analyzed  = analyze_doc(doc.content());
        let previous_version = self.docs.insert(doc.id.clone(), doc.clone());
        if let Some(previous_version) = previous_version {
            let previous_analyzed = analyze_doc(previous_version.content());
            for (ngram, highlights) in previous_analyzed {
                let is_empty = {
                    let docs_for_ngram = self.index.get_mut(&ngram).unwrap();
                    docs_for_ngram.remove(&IndexedDocument { doc: doc.clone(), highlighted: highlights });
                    docs_for_ngram.is_empty()
                };
                if is_empty {
                    self.index.remove(&ngram);
                }
            }
        }

        for (ngram, highlighted) in analyzed {
            self.index.entry(ngram).or_insert_with(|| HashSet::new())
                .insert(IndexedDocument { doc: doc.clone(), highlighted: highlighted });
        }
    }

    /// A basic search implementation that splits the query's content into whitespace-separated
    /// words, looks up the set of Documents for each word, and then concatenates the sets.
    pub fn search(&self, query: &str) -> Vec<SearchResult> {
        let unique_terms: HashSet<_> = query.split_whitespace().collect();
        let map = unique_terms.into_iter()
            .flat_map(|word| self.index.get(&word.to_lowercase()))
            .flat_map(|doc| doc)
            .cloned()
            .fold(HashMap::new(), |mut map, IndexedDocument { doc, highlighted: (begin, end) }| {
                match map.entry(doc) {
                    Vacant(entry) => {
                        let mut indices = HashMap::new();
                        indices.insert(begin, end);
                        entry.insert(indices);
                    }
                    Occupied(mut entry) => {
                        match entry.get_mut().entry(begin) {
                            Vacant(entry) => { entry.insert(end); }
                            Occupied(mut entry) => {
                                let current = entry.get_mut();
                                *current = std::cmp::max(*current, end);
                            }
                        }
                    }
                }
                map
            });
        let mut results: Vec<_> = map.into_iter()
            .map(|(doc, index_map)| SearchResult::new(doc, index_map.into_iter().collect()))
            .collect();
        results.sort_by(|result1, result2| result2.score.partial_cmp(&result1.score).unwrap());
        results
    }
}

fn analyze_doc(doc: &str)
-> iter::FlatMap<
    iter::Filter<
        GroupBy<bool, std::str::CharIndices, fn(&(usize, char)) -> bool>,
        fn(&(bool, Vec<(usize, char)>)) -> bool>,
    iter::Map<ops::Range<usize>, Ngrams>,
    fn((bool, Vec<(usize, char)>)) -> iter::Map<ops::Range<usize>, Ngrams>>
{
    doc.char_indices()
        .group_by(is_whitespace as fn(&(usize, char)) -> bool)
        .filter(not_whitespace as fn(&(bool, Vec<(usize, char)>)) -> bool)
        .flat_map(ngrams)
}

fn ngrams((_, chars): (bool, Vec<(usize, char)>)) -> iter::Map<ops::Range<usize>, Ngrams> {
    (1..chars.len() + 1).map(Ngrams::new(chars))
}

fn not_whitespace(&(is_whitespace, _): &(bool, Vec<(usize, char)>)) -> bool {
    !is_whitespace
}

fn is_whitespace(&(_, c): &(usize, char)) -> bool {
    c.is_whitespace()
}

struct Ngrams {
    chars: Vec<(usize, char)>
}

impl Ngrams {
    fn new(chars: Vec<(usize, char)>) -> Ngrams {
        Ngrams { chars: chars }
    }
}

impl Fn<(usize,)> for Ngrams {
    extern "rust-call" fn call(&self, (to,): (usize,)) -> (String, (usize, usize)) {
        let word = self.chars[..to].iter().flat_map(|&(_, c)| c.to_lowercase()).collect();
        let start = self.chars[0].0;
        let (last_idx, last_char) = self.chars[to - 1];
        let finish = last_idx + last_char.len_utf8();
        (word, (start, finish))
    }
}

impl FnMut<(usize,)> for Ngrams {
    extern "rust-call" fn call_mut(&mut self, to: (usize,)) -> (String, (usize, usize)) {
        self.call(to)
    }
}

impl FnOnce<(usize,)> for Ngrams {
    type Output = (String, (usize, usize));
    extern "rust-call" fn call_once(self, to: (usize,)) -> (String, (usize, usize)) {
        self.call(to)
    }
}

#[test]
fn test_search1() {
    let mut index = InvertedIndex::new();
    let doc1 = Document::new("1", "learn to program in rust today");
    index.index(doc1.clone());
    let doc2 = Document::new("2", "what did you today do");
    index.index(doc2.clone());
    let search_results = index.search("to");
    let expected = [
        SearchResult::new(Arc::new(doc1), vec![(6, 8), (25, 27)]),
        SearchResult::new(Arc::new(doc2), vec![(13, 15)])
    ];
    assert_eq!(search_results, expected.iter().cloned().collect::<Vec<_>>());
    assert_eq!("learn to program in rust today", expected[0].highlighted);
}

#[test]
fn test_ngrams() {
    let mut index = InvertedIndex::new();
    let doc1 = Document::new("1", "learn to program in rust today");
    let doc2 = Document::new("2", "what did you today do");
    index.index(doc1.clone());
    index.index(doc2.clone());
    let search_results = index.search("to");
    let expected = [
        SearchResult::new(Arc::new(doc1), vec![(6, 8), (25, 27)]),
        SearchResult::new(Arc::new(doc2), vec![(13, 15)]),
    ];
    assert_eq!(search_results, expected.iter().cloned().collect::<Vec<_>>());
    assert_eq!("learn to program in rust today", expected[0].highlighted);

}

#[test]
fn test_search2() {
    let mut index = InvertedIndex::new();
    let doc1 = Document::new("2", "what to do today");
    let doc2 = Document::new("3", "hey today");
    index.index(doc1.clone());
    index.index(doc2.clone());
    let search_results = index.search("to");
    let expected = [
        SearchResult::new(Arc::new(doc1), vec![(5, 7), (11, 13)]),
        SearchResult::new(Arc::new(doc2), vec![(4, 6)]),
    ];
    assert_eq!(search_results, expected.iter().cloned().collect::<Vec<_>>());
    assert_eq!(expected[0].highlighted, "what to do today");
}

#[test]
fn test_unicode() {
    let mut index = InvertedIndex::new();
    let doc = Document::new("0", "abc åäö");
    index.index(doc.clone());
    let to_search = "å";
    let search_results = index.search(to_search);
    let &SearchResult { ref doc, ref highlights, .. } = search_results.iter().next().unwrap();
    let (begin, end) = highlights[0];
    assert_eq!(&doc.content()[begin..end], to_search);
}

#[test]
fn test_update_doc() {
    let mut index = InvertedIndex::new();
    let doc = Document::new("0", "abc åäö");
    index.index(doc);
    let doc = Document::new("0", "different");
    index.index(doc);
    let search_results = index.search("å");
    assert!(search_results.is_empty());
    assert_eq!(index.docs.len(), 1);
}

#[test]
fn test_ranking() {
    let mut index = InvertedIndex::new();
    let doc = Document::new("0", "beat");
    index.index(doc.clone());
    let doc2 = Document::new("1", "beast");
    index.index(doc2.clone());
    let search_results = index.search("be");
    assert_eq!(index.docs.len(), 2);
    assert_eq!(&*search_results.into_iter().next().unwrap().doc, &doc);
}

#[test]
fn test_duplicate_term() {
    let mut index = InvertedIndex::new();
    let doc = Document::new("0", "beat");
    index.index(doc.clone());
    let search_results = index.search("be be");
    assert_eq!(search_results.len(), 1);
}

#[test]
fn test_duplicate_term2() {
    let mut index = InvertedIndex::new();
    let doc = Document::new("0", "beat");
    index.index(doc.clone());
    let search_results = index.search("be b");
    assert_eq!(search_results.len(), 1);
    assert_eq!(search_results[0], SearchResult::new(Arc::new(doc), vec![(0, 2)]));
}
<|endoftext|>


<commit_before>// Copyright 2014-2015 Sam Doshi (sam@metal-fish.co.uk)
// Copyright 2013-2014 Philippe Delrieu (philippe.delrieu@free.fr)
//
// Licensed under either of
//          Apache License, Version 2.0, (LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0)
//          MIT license (LICENSE-MIT or http://opensource.org/licenses/MIT).
// This file may not be copied, modified, or distributed except according to those terms.
use std::ptr;
use std::os::raw::c_char;

mod ffi;
use ffi::MaybeError;
pub mod types;
pub use types::*;

// Global fns
// ----------
/// `initialize` initalizes the underlying PortMidi C library, call this
/// before using the library.
///
/// Once initialized, PortMidi will no longer pickup any new Midi devices that are
/// connected, i.e. it does not support hot plugging.
pub fn initialize() -> PortMidiResult<()> {
    PortMidiResult::from(unsafe { ffi::Pm_Initialize() })
}

/// `terminate` terminates the underlying PortMidi C library, call this
/// after using the library.
pub fn terminate() -> PortMidiResult<()> {
    PortMidiResult::from(unsafe { ffi::Pm_Terminate() })
}

/// Return the number of devices. This number will not change during the lifetime
/// of the program.
pub fn count_devices() -> PortMidiDeviceId {
    unsafe { ffi::Pm_CountDevices() }
}

/// Gets the `PortMidiDeviceId` for the default input, or `None` if
/// there isn't one
///
/// See the PortMidi documentation for details of how to set the default device
pub fn get_default_input_device_id() -> Option<PortMidiDeviceId> {
    let id = unsafe { ffi::Pm_GetDefaultInputDeviceID() };
    if id == ffi::PM_NO_DEVICE {
        None
    } else {
        Some(id)
    }
}

/// Gets the `PortMidiDeviceId` for the default output, or `None` if
/// there isn't one
///
/// See the PortMidi documentation for details of how to set the default device
pub fn get_default_output_device_id() -> Option<PortMidiDeviceId> {
    let id = unsafe { ffi::Pm_GetDefaultOutputDeviceID() };
    if id == ffi::PM_NO_DEVICE {
        None
    } else {
        Some(id)
    }
}

// DeviceInfo
// ----------
/// Represents what we know about a device
#[derive(Clone, Debug)]
pub struct DeviceInfo {
    /// The `PortMidiDeviceId` used with `OutputPort::new` and `InputPort::new`
    pub device_id: PortMidiDeviceId,
    /// The name of the device
    pub name: String,
    /// Is the device an input
    pub input: bool,
    /// Is the device an output
    pub output: bool,
}

impl DeviceInfo {
    fn wrap(device_id: PortMidiDeviceId, device_info: *const ffi::PmDeviceInfo) -> DeviceInfo {
        let name = unsafe {
            let bytes = std::ffi::CStr::from_ptr((*device_info).name).to_bytes();
            std::str::from_utf8_unchecked(bytes).to_string()
        };
        let input = unsafe { (*device_info).input };
        let output = unsafe { (*device_info).output };

        DeviceInfo {
            device_id: device_id,
            name: name,
            input: input > 0,
            output: output > 0,
        }
    }
}

/// Returns a `DeviceInfo` with information about a device, or `None` if
/// it does not exist
pub fn get_device_info(device_id: PortMidiDeviceId) -> Option<DeviceInfo> {
    let info = unsafe { ffi::Pm_GetDeviceInfo(device_id) };
    if info.is_null() {
        None
    } else {
        Some(DeviceInfo::wrap(device_id, info))
    }
}

// Midi events
// -----------
/// Represents a single midi message, see also `MidiEvent`
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct MidiMessage {
    pub status: u8,
    pub data1: u8,
    pub data2: u8,
}

impl MidiMessage {
    fn wrap(cmessage: ffi::PmMessage) -> MidiMessage {
        MidiMessage {
            status: ((cmessage) & 0xFF) as u8,
            data1: (((cmessage) >> 8) & 0xFF) as u8,
            data2: (((cmessage) >> 16) & 0xFF) as u8,
        }
    }

    fn unwrap(&self) -> ffi::PmMessage {
        ((((self.data2 as i32) << 16) & 0xFF0000) | (((self.data1 as i32) << 8) & 0xFF00) |
         ((self.status as i32) & 0xFF)) as i32
    }
}

/// Represents a time stamped midi event. See also `MidiMessage`
///
/// See the PortMidi documentation for how SysEx and midi realtime messages
/// are handled
///
/// TODO: what to do about the timestamp?
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct MidiEvent {
    pub message: MidiMessage,
    pub timestamp: ffi::PmTimestamp,
}

impl MidiEvent {
    fn wrap(event: ffi::PmEvent) -> MidiEvent {
        MidiEvent {
            message: MidiMessage::wrap(event.message),
            timestamp: event.timestamp,
        }
    }

    fn unwrap(&self) -> ffi::PmEvent {
        ffi::PmEvent {
            message: self.message.unwrap(),
            timestamp: self.timestamp,
        }
    }
}

// Input
// -----
/// Representation of an input midi port
#[allow(missing_copy_implementations)]
pub struct InputPort {
    stream: *const ffi::PortMidiStream,
    input_device: ffi::PmDeviceId,
    buffer_size: i32,
}

const EVENT_BUFFER_SIZE: usize = 128;
impl InputPort {
    /// Construct a new `InputPort` for `input_device`
    pub fn new(input_device: PortMidiDeviceId, buffer_size: i32) -> InputPort {
        InputPort {
            stream: ptr::null(),
            input_device: input_device,
            buffer_size: buffer_size,
        }
    }

    /// Open the port returning an error if there is a problem
    pub fn open(&mut self) -> PortMidiResult<()> {
        PortMidiResult::from(unsafe {
            ffi::Pm_OpenInput(&self.stream,
                              self.input_device,
                              ptr::null(),
                              self.buffer_size,
                              ptr::null(),
                              ptr::null())
        })
    }

    pub fn read_n(&mut self, cnt: usize) -> PortMidiResult<Option<Vec<MidiEvent>>> {
        let read_cnt = if cnt > EVENT_BUFFER_SIZE {
            EVENT_BUFFER_SIZE as i32
        } else {
            cnt as i32
        };
        let mut event_buffer = [ffi::PmEvent::default(); EVENT_BUFFER_SIZE];
        let res = unsafe {
            ffi::Pm_Read(self.stream,
                         event_buffer.as_mut_ptr(),
                         read_cnt)
        };
        if res < 0 {
            let err = ffi::PmError::try_from(res).unwrap();
            // TODO: Return the error
            println!("error: {:?}", err);
            return Ok(None);
        } else if res == 0 {
            Ok(None)
        } else {
            // remove mutability and replace return value
            let mut events = (0..res as usize)
                             .map(|i| MidiEvent::wrap(event_buffer[i].clone()))
                             .collect::<Vec<MidiEvent>>();
            Ok(Some(events))
        }
    }

    /// Reads a single `MidiEvent` if one is avaible
    ///
    /// A `Result` of `None` means no event was available.
    ///
    /// See the PortMidi documentation for information on how it deals with input
    /// overflows
    /// TODO: call `read_n`
    pub fn read(&mut self) -> PortMidiResult<Option<MidiEvent>> {
        match self.read_n(1) {
            Ok(Some(mut vec)) => Ok(vec.pop()),
            Ok(_) => Ok(None),
            Err(e) => Err(e),
        }
    }

    /// `poll` tests if there is input available, either returing a bool or an error
    pub fn poll(&self) -> PortMidiResult<bool> {
        let pm_error = unsafe { ffi::Pm_Poll(self.stream) };
        match pm_error {
            ffi::PmError::PmNoError => Ok(false),
            ffi::PmError::PmGotData => Ok(true),
            other => PortMidiResult::from(other).map(|_| false),
        }
    }

    /// Closes the input, flushing any pending buffers
    ///
    /// PortMidi attempts to close open streams when the application
    /// exits, but this can be difficult under Windows
    /// (according to the PortMidi documentation).
    pub fn close(&mut self) -> PortMidiResult<()> {
        PortMidiResult::from(unsafe { ffi::Pm_Close(self.stream) })
    }

    //    Test whether stream has a pending host error. Normally, the client finds
    //    out about errors through returned error codes, but some errors can occur
    //    asynchronously where the client does not
    //    explicitly call a function, and therefore cannot receive an error code.
    //    The client can test for a pending error using has_host_error(). If true,
    //    the error can be accessed and cleared by calling get_Error_text().
    //    Errors are also cleared by calling other functions that can return
    //    errors, e.g. open_input(), open_output(), read(), write(). The
    //    client does not need to call Pm_HasHostError(). Any pending error will be
    //    reported the next time the client performs an explicit function call on
    //    the stream, e.g. an input or output operation. Until the error is cleared,
    //    no new error codes will be obtained, even for a different stream.
    //
    fn has_host_error(&self) -> bool {
        unsafe { ffi::Pm_HasHostError(self.stream) > 0 }
    }
}

// Output
// ------
/// Representation of an output midi port
#[allow(missing_copy_implementations)]
pub struct OutputPort {
    stream: *const ffi::PortMidiStream,
    output_device: ffi::PmDeviceId,
    buffer_size: i32,
}
impl OutputPort {
    /// Construct a new `InputPort` for `input_device`
    pub fn new(output_device: PortMidiDeviceId, buffer_size: i32) -> OutputPort {
        OutputPort {
            stream: ptr::null(),
            output_device: output_device,
            buffer_size: buffer_size,
        }
    }

    /// Open the port returning an error if there is a problem
    pub fn open(&mut self) -> PortMidiResult<()> {
        PortMidiResult::from(unsafe {
            ffi::Pm_OpenOutput(&self.stream,
                               self.output_device,
                               ptr::null(),
                               self.buffer_size,
                               ptr::null(),
                               ptr::null(),
                               0)
        })
    }

    /// Terminates outgoing messages immediately
    ///
    /// The caller should immediately close the output port, this may
    /// result in transmission of a partial midi message. Note, not all platforms
    /// support abort.
    pub fn abort(&mut self) -> PortMidiResult<()> {
        PortMidiResult::from(unsafe { ffi::Pm_Abort(self.stream) })
    }

    /// Closes the midi stream, flushing any pending buffers
    ///
    /// PortMidi attempts to close open streams when the application
    /// exits, but this can be difficult under Windows
    /// (according to the PortMidi documentation).
    pub fn close(&mut self) -> PortMidiResult<()> {
        PortMidiResult::from(unsafe { ffi::Pm_Close(self.stream) })
    }

    /// Write a single `MidiEvent`
    pub fn write_event(&mut self, midi_event: MidiEvent) -> PortMidiResult<()> {
        let event = midi_event.unwrap();
        PortMidiResult::from(unsafe { ffi::Pm_Write(self.stream, &event, 1) })
    }

    /// Write a single `MidiMessage` immediately
    pub fn write_message(&mut self, midi_message: MidiMessage) -> PortMidiResult<()> {
        let message = midi_message.unwrap();
        PortMidiResult::from(unsafe { ffi::Pm_WriteShort(self.stream, 0, message) })
    }

    //    Test whether stream has a pending host error. Normally, the client finds
    //    out about errors through returned error codes, but some errors can occur
    //    asynchronously where the client does not
    //    explicitly call a function, and therefore cannot receive an error code.
    //    The client can test for a pending error using has_host_error(). If true,
    //    the error can be accessed and cleared by calling get_Error_text().
    //    Errors are also cleared by calling other functions that can return
    //    errors, e.g. open_input(), open_output(), read(), write(). The
    //    client does not need to call Pm_HasHostError(). Any pending error will be
    //    reported the next time the client performs an explicit function call on
    //    the stream, e.g. an input or output operation. Until the error is cleared,
    //    no new error codes will be obtained, even for a different stream.
    //
    pub fn has_host_error(&self) -> bool {
        unsafe { ffi::Pm_HasHostError(self.stream) > 0 }
    }
}

// Old code
// --------
/**  Translate portmidi error number into human readable message.
*    These strings are constants (set at compile time) so client has
*    no need to allocate storage
*/
pub fn get_error_text(error_code: ffi::PmError) -> String {
    unsafe {
        let error_text = ffi::Pm_GetErrorText(error_code);
        let bytes = std::ffi::CStr::from_ptr(error_text).to_bytes();
        std::str::from_utf8_unchecked(bytes).to_string()
    }
}

/**  Translate portmidi host error into human readable message.
     These strings are computed at run time, so client has to allocate storage.
     After this routine executes, the host error is cleared.
*/
pub fn get_host_error_text(msg: *const c_char, len: i32) {
    unsafe {
        ffi::Pm_GetHostErrorText(msg, len);
    }
}

pub const HDRLENGTH: i32 = 50;
pub const PM_HOST_ERROR_MSG_LEN: i32 = 256;
<commit_msg>Remove `has_host_error`.<commit_after>// Copyright 2014-2015 Sam Doshi (sam@metal-fish.co.uk)
// Copyright 2013-2014 Philippe Delrieu (philippe.delrieu@free.fr)
//
// Licensed under either of
//          Apache License, Version 2.0, (LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0)
//          MIT license (LICENSE-MIT or http://opensource.org/licenses/MIT).
// This file may not be copied, modified, or distributed except according to those terms.
use std::ptr;
use std::os::raw::c_char;

mod ffi;
use ffi::MaybeError;
pub mod types;
pub use types::*;

// Global fns
// ----------
/// `initialize` initalizes the underlying PortMidi C library, call this
/// before using the library.
///
/// Once initialized, PortMidi will no longer pickup any new Midi devices that are
/// connected, i.e. it does not support hot plugging.
pub fn initialize() -> PortMidiResult<()> {
    PortMidiResult::from(unsafe { ffi::Pm_Initialize() })
}

/// `terminate` terminates the underlying PortMidi C library, call this
/// after using the library.
pub fn terminate() -> PortMidiResult<()> {
    PortMidiResult::from(unsafe { ffi::Pm_Terminate() })
}

/// Return the number of devices. This number will not change during the lifetime
/// of the program.
pub fn count_devices() -> PortMidiDeviceId {
    unsafe { ffi::Pm_CountDevices() }
}

/// Gets the `PortMidiDeviceId` for the default input, or `None` if
/// there isn't one
///
/// See the PortMidi documentation for details of how to set the default device
pub fn get_default_input_device_id() -> Option<PortMidiDeviceId> {
    let id = unsafe { ffi::Pm_GetDefaultInputDeviceID() };
    if id == ffi::PM_NO_DEVICE {
        None
    } else {
        Some(id)
    }
}

/// Gets the `PortMidiDeviceId` for the default output, or `None` if
/// there isn't one
///
/// See the PortMidi documentation for details of how to set the default device
pub fn get_default_output_device_id() -> Option<PortMidiDeviceId> {
    let id = unsafe { ffi::Pm_GetDefaultOutputDeviceID() };
    if id == ffi::PM_NO_DEVICE {
        None
    } else {
        Some(id)
    }
}

// DeviceInfo
// ----------
/// Represents what we know about a device
#[derive(Clone, Debug)]
pub struct DeviceInfo {
    /// The `PortMidiDeviceId` used with `OutputPort::new` and `InputPort::new`
    pub device_id: PortMidiDeviceId,
    /// The name of the device
    pub name: String,
    /// Is the device an input
    pub input: bool,
    /// Is the device an output
    pub output: bool,
}

impl DeviceInfo {
    fn wrap(device_id: PortMidiDeviceId, device_info: *const ffi::PmDeviceInfo) -> DeviceInfo {
        let name = unsafe {
            let bytes = std::ffi::CStr::from_ptr((*device_info).name).to_bytes();
            std::str::from_utf8_unchecked(bytes).to_string()
        };
        let input = unsafe { (*device_info).input };
        let output = unsafe { (*device_info).output };

        DeviceInfo {
            device_id: device_id,
            name: name,
            input: input > 0,
            output: output > 0,
        }
    }
}

/// Returns a `DeviceInfo` with information about a device, or `None` if
/// it does not exist
pub fn get_device_info(device_id: PortMidiDeviceId) -> Option<DeviceInfo> {
    let info = unsafe { ffi::Pm_GetDeviceInfo(device_id) };
    if info.is_null() {
        None
    } else {
        Some(DeviceInfo::wrap(device_id, info))
    }
}

// Midi events
// -----------
/// Represents a single midi message, see also `MidiEvent`
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct MidiMessage {
    pub status: u8,
    pub data1: u8,
    pub data2: u8,
}

impl MidiMessage {
    fn wrap(cmessage: ffi::PmMessage) -> MidiMessage {
        MidiMessage {
            status: ((cmessage) & 0xFF) as u8,
            data1: (((cmessage) >> 8) & 0xFF) as u8,
            data2: (((cmessage) >> 16) & 0xFF) as u8,
        }
    }

    fn unwrap(&self) -> ffi::PmMessage {
        ((((self.data2 as i32) << 16) & 0xFF0000) | (((self.data1 as i32) << 8) & 0xFF00) |
         ((self.status as i32) & 0xFF)) as i32
    }
}

/// Represents a time stamped midi event. See also `MidiMessage`
///
/// See the PortMidi documentation for how SysEx and midi realtime messages
/// are handled
///
/// TODO: what to do about the timestamp?
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct MidiEvent {
    pub message: MidiMessage,
    pub timestamp: ffi::PmTimestamp,
}

impl MidiEvent {
    fn wrap(event: ffi::PmEvent) -> MidiEvent {
        MidiEvent {
            message: MidiMessage::wrap(event.message),
            timestamp: event.timestamp,
        }
    }

    fn unwrap(&self) -> ffi::PmEvent {
        ffi::PmEvent {
            message: self.message.unwrap(),
            timestamp: self.timestamp,
        }
    }
}

// Input
// -----
/// Representation of an input midi port
#[allow(missing_copy_implementations)]
pub struct InputPort {
    stream: *const ffi::PortMidiStream,
    input_device: ffi::PmDeviceId,
    buffer_size: i32,
}

const EVENT_BUFFER_SIZE: usize = 128;
impl InputPort {
    /// Construct a new `InputPort` for `input_device`
    pub fn new(input_device: PortMidiDeviceId, buffer_size: i32) -> InputPort {
        InputPort {
            stream: ptr::null(),
            input_device: input_device,
            buffer_size: buffer_size,
        }
    }

    /// Open the port returning an error if there is a problem
    pub fn open(&mut self) -> PortMidiResult<()> {
        PortMidiResult::from(unsafe {
            ffi::Pm_OpenInput(&self.stream,
                              self.input_device,
                              ptr::null(),
                              self.buffer_size,
                              ptr::null(),
                              ptr::null())
        })
    }

    pub fn read_n(&mut self, cnt: usize) -> PortMidiResult<Option<Vec<MidiEvent>>> {
        let read_cnt = if cnt > EVENT_BUFFER_SIZE {
            EVENT_BUFFER_SIZE as i32
        } else {
            cnt as i32
        };
        let mut event_buffer = [ffi::PmEvent::default(); EVENT_BUFFER_SIZE];
        let res = unsafe {
            ffi::Pm_Read(self.stream,
                         event_buffer.as_mut_ptr(),
                         read_cnt)
        };
        if res < 0 {
            let err = ffi::PmError::try_from(res).unwrap();
            // TODO: Return the error
            println!("error: {:?}", err);
            return Ok(None);
        } else if res == 0 {
            Ok(None)
        } else {
            // remove mutability and replace return value
            let mut events = (0..res as usize)
                             .map(|i| MidiEvent::wrap(event_buffer[i].clone()))
                             .collect::<Vec<MidiEvent>>();
            Ok(Some(events))
        }
    }

    /// Reads a single `MidiEvent` if one is avaible
    ///
    /// A `Result` of `None` means no event was available.
    ///
    /// See the PortMidi documentation for information on how it deals with input
    /// overflows
    /// TODO: call `read_n`
    pub fn read(&mut self) -> PortMidiResult<Option<MidiEvent>> {
        match self.read_n(1) {
            Ok(Some(mut vec)) => Ok(vec.pop()),
            Ok(_) => Ok(None),
            Err(e) => Err(e),
        }
    }

    /// `poll` tests if there is input available, either returing a bool or an error
    pub fn poll(&self) -> PortMidiResult<bool> {
        let pm_error = unsafe { ffi::Pm_Poll(self.stream) };
        match pm_error {
            ffi::PmError::PmNoError => Ok(false),
            ffi::PmError::PmGotData => Ok(true),
            other => PortMidiResult::from(other).map(|_| false),
        }
    }

    /// Closes the input, flushing any pending buffers
    ///
    /// PortMidi attempts to close open streams when the application
    /// exits, but this can be difficult under Windows
    /// (according to the PortMidi documentation).
    pub fn close(&mut self) -> PortMidiResult<()> {
        PortMidiResult::from(unsafe { ffi::Pm_Close(self.stream) })
    }
}

// Output
// ------
/// Representation of an output midi port
#[allow(missing_copy_implementations)]
pub struct OutputPort {
    stream: *const ffi::PortMidiStream,
    output_device: ffi::PmDeviceId,
    buffer_size: i32,
}
impl OutputPort {
    /// Construct a new `InputPort` for `input_device`
    pub fn new(output_device: PortMidiDeviceId, buffer_size: i32) -> OutputPort {
        OutputPort {
            stream: ptr::null(),
            output_device: output_device,
            buffer_size: buffer_size,
        }
    }

    /// Open the port returning an error if there is a problem
    pub fn open(&mut self) -> PortMidiResult<()> {
        PortMidiResult::from(unsafe {
            ffi::Pm_OpenOutput(&self.stream,
                               self.output_device,
                               ptr::null(),
                               self.buffer_size,
                               ptr::null(),
                               ptr::null(),
                               0)
        })
    }

    /// Terminates outgoing messages immediately
    ///
    /// The caller should immediately close the output port, this may
    /// result in transmission of a partial midi message. Note, not all platforms
    /// support abort.
    pub fn abort(&mut self) -> PortMidiResult<()> {
        PortMidiResult::from(unsafe { ffi::Pm_Abort(self.stream) })
    }

    /// Closes the midi stream, flushing any pending buffers
    ///
    /// PortMidi attempts to close open streams when the application
    /// exits, but this can be difficult under Windows
    /// (according to the PortMidi documentation).
    pub fn close(&mut self) -> PortMidiResult<()> {
        PortMidiResult::from(unsafe { ffi::Pm_Close(self.stream) })
    }

    /// Write a single `MidiEvent`
    pub fn write_event(&mut self, midi_event: MidiEvent) -> PortMidiResult<()> {
        let event = midi_event.unwrap();
        PortMidiResult::from(unsafe { ffi::Pm_Write(self.stream, &event, 1) })
    }

    /// Write a single `MidiMessage` immediately
    pub fn write_message(&mut self, midi_message: MidiMessage) -> PortMidiResult<()> {
        let message = midi_message.unwrap();
        PortMidiResult::from(unsafe { ffi::Pm_WriteShort(self.stream, 0, message) })
    }
}

// Old code
// --------
/**  Translate portmidi error number into human readable message.
*    These strings are constants (set at compile time) so client has
*    no need to allocate storage
*/
pub fn get_error_text(error_code: ffi::PmError) -> String {
    unsafe {
        let error_text = ffi::Pm_GetErrorText(error_code);
        let bytes = std::ffi::CStr::from_ptr(error_text).to_bytes();
        std::str::from_utf8_unchecked(bytes).to_string()
    }
}

/**  Translate portmidi host error into human readable message.
     These strings are computed at run time, so client has to allocate storage.
     After this routine executes, the host error is cleared.
*/
pub fn get_host_error_text(msg: *const c_char, len: i32) {
    unsafe {
        ffi::Pm_GetHostErrorText(msg, len);
    }
}

pub const HDRLENGTH: i32 = 50;
pub const PM_HOST_ERROR_MSG_LEN: i32 = 256;
<|endoftext|>