<commit_before>use std::cmp::Eq;
use std::fmt::Display;
use std::iter::FromIterator;

trait Parser<T: Display + Eq, E>: Sized {
    fn preview(&self) -> Option<&T>;
    fn consume(&mut self) -> Option<T>;
    fn current_pos(&self) -> (i32, i32);
    fn error<S: Into<String>>(&self, message: S) -> E;

    fn export(&self) -> Self;
    fn import(&mut self, backup: Self);

    fn next(&mut self) -> Result<T, E> {
        self.consume().ok_or(self.error("unexpected eof"))
    }

    fn predicate<F>(&mut self, pred: F) -> Result<T, E>
        where F: Fn(&T) -> bool
    {
        let x = try!(self.next());
        if pred(&x) {
            Ok(x)
        } else {
            Err(self.error(format!("unexpected token {}", x)))
        }
    }

    fn atom(&mut self, expected: T) -> Result<T, E> {
        let x = try!(self.next());
        if x == expected {
            Ok(x)
        } else {
            Err(self.error(format!("unexpected token {}, expected {}", x, expected)))
        }
    }

    fn string<S, O>(&mut self, s: S) -> Result<O, E>
        where S: IntoIterator<Item = T>,
              O: FromIterator<T>
    {
        let mut res: Vec<T> = Vec::new();
        for c in s {
            match self.atom(c) {
                Ok(x) => res.push(x),
                Err(err) => return Err(err),
            }
        }
        Ok(O::from_iter(res))
    }

    fn try<O, F>(&mut self, parser: F) -> Result<O, E>
        where F: Fn(&mut Self) -> Result<O, E>
    {
        let backup = self.export();
        parser(self).map_err(|x| {
                                 self.import(backup);
                                 x
                             })
    }

    fn choose<O, F>(&mut self, parsers: &[F]) -> Result<O, E>
        where F: Fn(&mut Self) -> Result<O, E>
    {
        for parser in parsers {
            match self.try(|p| parser(p)) {
                Ok(x) => return Ok(x),
                Err(_) => continue,
            }
        }

        Err(self.error(match self.preview() {
                           Some(x) => format!("unexpected token {}", x),
                           None => String::from("unexpected eof"),
                       }))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::vec_deque::VecDeque;

    struct TP {
        input: VecDeque<i32>,
    }

    impl TP {
        fn new(input: &[i32]) -> TP {
            TP { input: VecDeque::from(input.to_vec()) }
        }
    }

    impl Parser<i32, String> for TP {
        fn consume(&mut self) -> Option<i32> {
            self.input.pop_front()
        }

        fn preview(&self) -> Option<&i32> {
            self.input.front()
        }

        fn current_pos(&self) -> (i32, i32) {
            (0, 0)
        }

        fn error<S: Into<String>>(&self, message: S) -> String {
            message.into()
        }

        fn export(&self) -> Self {
            TP { input: self.input.clone() }
        }

        fn import(&mut self, backup: Self) {
            self.input = backup.input;
        }
    }

    type TPR = Result<Vec<i32>, String>;

    fn err<T>(m: &str) -> Result<T, String> {
        Err(String::from(m))
    }

    #[test]
    fn next_success() {
        let mut p = TP::new(&[1, 2, 3]);
        assert_eq!(p.next(), Ok(1));
        assert_eq!(p.next(), Ok(2));
        assert_eq!(p.next(), Ok(3));
    }

    #[test]
    fn next_fail_empty() {
        let mut p = TP::new(&[]);
        assert_eq!(p.next(), err("unexpected eof"));
    }

    #[test]
    fn predicate_success() {
        let mut p = TP::new(&[2, 4, 6]);
        assert_eq!(p.predicate(|x| x % 2 == 0), Ok(2));
        assert_eq!(p.predicate(|x| x % 2 == 0), Ok(4));
        assert_eq!(p.predicate(|x| x % 2 == 0), Ok(6));
    }

    #[test]
    fn predicate_fail_empty() {
        let mut p = TP::new(&[]);
        assert_eq!(p.predicate(|x| x % 2 == 0), err("unexpected eof"));
    }

    #[test]
    fn predicate_fail_not_satisfy() {
        let mut p = TP::new(&[3, 5, 7]);
        assert_eq!(p.predicate(|x| x % 2 == 0), err("unexpected token 3"));
    }

    #[test]
    fn atom_success() {
        let mut p = TP::new(&[2, 4, 6]);
        assert_eq!(p.atom(2), Ok(2));
        assert_eq!(p.atom(4), Ok(4));
        assert_eq!(p.atom(6), Ok(6));
    }

    #[test]
    fn atom_fail_empty() {
        let mut p = TP::new(&[]);
        assert_eq!(p.atom(2), err("unexpected eof"));
    }

    #[test]
    fn atom_fail_not_expected() {
        let mut p = TP::new(&[3, 5, 7]);
        assert_eq!(p.atom(3), Ok(3));
        assert_eq!(p.atom(4), err("unexpected token 5, expected 4"));
    }

    #[test]
    fn string_success() {
        let mut p = TP::new(&[2, 4, 6]);
        assert_eq!(p.string(vec![2, 4, 6]), Ok(vec![2, 4, 6]));
    }

    #[test]
    fn string_fail_empty() {
        let mut p = TP::new(&[]);
        assert_eq!(p.string(vec![2, 4, 6]) as TPR, err("unexpected eof"));
    }

    #[test]
    fn string_fail_not_expected() {
        let mut p = TP::new(&[2, 5, 6]);
        assert_eq!(p.string(vec![2, 4, 6]) as TPR,
                   err("unexpected token 5, expected 4"));
    }

    #[test]
    fn try_success() {
        let mut p = TP::new(&[2, 4, 6]);
        assert_eq!(p.try(|p| p.string(vec![2, 4, 6])), Ok(vec![2, 4, 6]));
    }

    #[test]
    fn try_fail_recover() {
        let mut p = TP::new(&[2, 4, 6]);
        assert_eq!(p.try(|p| p.string(vec![2, 4, 7])) as TPR,
                   err("unexpected token 6, expected 7"));
        assert_eq!(p.try(|p| p.string(vec![2, 4, 6])), Ok(vec![2, 4, 6]));
    }

    #[test]
    fn choose_success() {
        let mut p = TP::new(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
        assert_eq!(p.choose(&[|p| p.string(vec![1, 2, 3]),
                              |p| p.string(vec![4, 5, 6, 7]),
                              |p| p.string(vec![4, 5, 6])]),
                   Ok(vec![1, 2, 3]));
    }

    #[test]
    fn choose_success_with_recover() {
        let mut p = TP::new(&[4, 5, 6, 7, 8, 9, 10]);
        assert_eq!(p.choose(&[|p| p.string(vec![1, 2, 3]),
                              |p| p.string(vec![4, 5, 6, 8]),
                              |p| p.string(vec![4, 5, 6])]),
                   Ok(vec![4, 5, 6]));
    }

    #[test]
    fn fail_choose_no_match() {
        let mut p = TP::new(&[5, 6, 7, 8, 9, 10]);
        assert_eq!(p.choose(&[|p| p.string(vec![1, 2, 3]),
                              |p| p.string(vec![4, 5, 6, 8]),
                              |p| p.string(vec![4, 5, 6])]),
                   err("unexpected token 5"));
    }

    #[test]
    fn fail_choose_empty() {
        let mut p = TP::new(&[]);
        assert_eq!(p.choose(&[|p| p.string(vec![1, 2, 3]),
                              |p| p.string(vec![4, 5, 6, 8]),
                              |p| p.string(vec![4, 5, 6])]),
                   err("unexpected eof"));
    }
}
<commit_msg>Use trait object for choose combinator<commit_after>use std::cmp::Eq;
use std::fmt::Display;
use std::iter::FromIterator;

trait Parser<T: Display + Eq, E>: Sized {
    fn preview(&self) -> Option<&T>;
    fn consume(&mut self) -> Option<T>;
    fn current_pos(&self) -> (i32, i32);
    fn error<S: Into<String>>(&self, message: S) -> E;

    fn export(&self) -> Self;
    fn import(&mut self, backup: Self);

    fn next(&mut self) -> Result<T, E> {
        self.consume().ok_or(self.error("unexpected eof"))
    }

    fn predicate<F>(&mut self, pred: F) -> Result<T, E>
        where F: Fn(&T) -> bool
    {
        let x = try!(self.next());
        if pred(&x) {
            Ok(x)
        } else {
            Err(self.error(format!("unexpected token {}", x)))
        }
    }

    fn atom(&mut self, expected: T) -> Result<T, E> {
        let x = try!(self.next());
        if x == expected {
            Ok(x)
        } else {
            Err(self.error(format!("unexpected token {}, expected {}", x, expected)))
        }
    }

    fn string<S, O>(&mut self, s: S) -> Result<O, E>
        where S: IntoIterator<Item = T>,
              O: FromIterator<T>
    {
        let mut res: Vec<T> = Vec::new();
        for c in s {
            match self.atom(c) {
                Ok(x) => res.push(x),
                Err(err) => return Err(err),
            }
        }
        Ok(O::from_iter(res))
    }

    fn try<O, F>(&mut self, parser: F) -> Result<O, E>
        where F: Fn(&mut Self) -> Result<O, E>
    {
        let backup = self.export();
        parser(self).map_err(|x| {
                                 self.import(backup);
                                 x
                             })
    }

    fn choose<O>(&mut self, parsers: &[&Fn(&mut Self) -> Result<O, E>]) -> Result<O, E> {
        for parser in parsers {
            match self.try(|p| parser(p)) {
                Ok(x) => return Ok(x),
                Err(_) => continue,
            }
        }

        Err(self.error(match self.preview() {
                           Some(x) => format!("unexpected token {}", x),
                           None => String::from("unexpected eof"),
                       }))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::vec_deque::VecDeque;

    struct TP {
        input: VecDeque<i32>,
    }

    impl TP {
        fn new(input: &[i32]) -> TP {
            TP { input: VecDeque::from(input.to_vec()) }
        }
    }

    impl Parser<i32, String> for TP {
        fn consume(&mut self) -> Option<i32> {
            self.input.pop_front()
        }

        fn preview(&self) -> Option<&i32> {
            self.input.front()
        }

        fn current_pos(&self) -> (i32, i32) {
            (0, 0)
        }

        fn error<S: Into<String>>(&self, message: S) -> String {
            message.into()
        }

        fn export(&self) -> Self {
            TP { input: self.input.clone() }
        }

        fn import(&mut self, backup: Self) {
            self.input = backup.input;
        }
    }

    type TPR = Result<Vec<i32>, String>;

    fn err<T>(m: &str) -> Result<T, String> {
        Err(String::from(m))
    }

    #[test]
    fn next_success() {
        let mut p = TP::new(&[1, 2, 3]);
        assert_eq!(p.next(), Ok(1));
        assert_eq!(p.next(), Ok(2));
        assert_eq!(p.next(), Ok(3));
    }

    #[test]
    fn next_fail_empty() {
        let mut p = TP::new(&[]);
        assert_eq!(p.next(), err("unexpected eof"));
    }

    #[test]
    fn predicate_success() {
        let mut p = TP::new(&[2, 4, 6]);
        assert_eq!(p.predicate(|x| x % 2 == 0), Ok(2));
        assert_eq!(p.predicate(|x| x % 2 == 0), Ok(4));
        assert_eq!(p.predicate(|x| x % 2 == 0), Ok(6));
    }

    #[test]
    fn predicate_fail_empty() {
        let mut p = TP::new(&[]);
        assert_eq!(p.predicate(|x| x % 2 == 0), err("unexpected eof"));
    }

    #[test]
    fn predicate_fail_not_satisfy() {
        let mut p = TP::new(&[3, 5, 7]);
        assert_eq!(p.predicate(|x| x % 2 == 0), err("unexpected token 3"));
    }

    #[test]
    fn atom_success() {
        let mut p = TP::new(&[2, 4, 6]);
        assert_eq!(p.atom(2), Ok(2));
        assert_eq!(p.atom(4), Ok(4));
        assert_eq!(p.atom(6), Ok(6));
    }

    #[test]
    fn atom_fail_empty() {
        let mut p = TP::new(&[]);
        assert_eq!(p.atom(2), err("unexpected eof"));
    }

    #[test]
    fn atom_fail_not_expected() {
        let mut p = TP::new(&[3, 5, 7]);
        assert_eq!(p.atom(3), Ok(3));
        assert_eq!(p.atom(4), err("unexpected token 5, expected 4"));
    }

    #[test]
    fn string_success() {
        let mut p = TP::new(&[2, 4, 6]);
        assert_eq!(p.string(vec![2, 4, 6]), Ok(vec![2, 4, 6]));
    }

    #[test]
    fn string_fail_empty() {
        let mut p = TP::new(&[]);
        assert_eq!(p.string(vec![2, 4, 6]) as TPR, err("unexpected eof"));
    }

    #[test]
    fn string_fail_not_expected() {
        let mut p = TP::new(&[2, 5, 6]);
        assert_eq!(p.string(vec![2, 4, 6]) as TPR,
                   err("unexpected token 5, expected 4"));
    }

    #[test]
    fn try_success() {
        let mut p = TP::new(&[2, 4, 6]);
        assert_eq!(p.try(|p| p.string(vec![2, 4, 6])), Ok(vec![2, 4, 6]));
    }

    #[test]
    fn try_fail_recover() {
        let mut p = TP::new(&[2, 4, 6]);
        assert_eq!(p.try(|p| p.string(vec![2, 4, 7])) as TPR,
                   err("unexpected token 6, expected 7"));
        assert_eq!(p.try(|p| p.string(vec![2, 4, 6])), Ok(vec![2, 4, 6]));
    }

    #[test]
    fn choose_success() {
        let mut p = TP::new(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
        assert_eq!(p.choose(&[&|p| p.string(vec![1, 2, 3]),
                              &|p| p.string(vec![4, 5, 6, 7]),
                              &|p| p.string(vec![4, 5, 6])]),
                   Ok(vec![1, 2, 3]));
    }

    #[test]
    fn choose_success_with_recover() {
        let mut p = TP::new(&[4, 5, 6, 7, 8, 9, 10]);
        assert_eq!(p.choose(&[&|p| p.string(vec![1, 2, 3]),
                              &|p| p.string(vec![4, 5, 6, 8]),
                              &|p| p.string(vec![4, 5, 6])]),
                   Ok(vec![4, 5, 6]));
    }

    #[test]
    fn fail_choose_no_match() {
        let mut p = TP::new(&[5, 6, 7, 8, 9, 10]);
        assert_eq!(p.choose(&[&|p| p.string(vec![1, 2, 3]),
                              &|p| p.string(vec![4, 5, 6, 8]),
                              &|p| p.string(vec![4, 5, 6])]) as TPR,
                   err("unexpected token 5"));
    }

    #[test]
    fn fail_choose_empty() {
        let mut p = TP::new(&[]);
        assert_eq!(p.choose(&[&|p| p.string(vec![1, 2, 3]),
                              &|p| p.string(vec![4, 5, 6, 8]),
                              &|p| p.string(vec![4, 5, 6])]) as TPR,
                   err("unexpected eof"));
    }
}
<|endoftext|>


<commit_before>use job_handle::ChildProcess;
use process_manager::PROCESS_MANAGER;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::RwLock;
use libc;
use libc::c_int;
use result::ShellResultExt;
use std::thread::ThreadId;
use std::thread;
use std::thread::JoinHandle;
use std::any::Any;
use std::cell::RefCell;

/// Thread local shell.
pub struct LocalShell {
    thread_id: ThreadId,
    processes: Vec<Arc<RwLock<ChildProcess>>>,
    signaled: bool
}

impl LocalShell {
    fn new() -> LocalShell {
        LocalShell {
            thread_id: thread::current().id(),
            processes: Vec::new(),
            signaled: false
        }
    }

    pub fn add_process(&mut self, process: &Arc<RwLock<ChildProcess>>) {
        self.processes.push(process.clone());
    }

    pub fn remove_process(&mut self, process: &Arc<RwLock<ChildProcess>>) {
        self.processes.retain(|p| !Arc::ptr_eq(p, process));
    }

    pub fn signal(&mut self, signal: c_int) {
        self.signaled = true;
        for process in &self.processes {
            let lock = process.read().unwrap();
            lock.signal(signal).print_error();
        }
    }

    pub fn signaled(&self) -> bool {
        self.signaled
    }

    pub fn thread_id(&self) -> &ThreadId {
        &self.thread_id
    }
}

struct LocalShellScope(ThreadId, Arc<Mutex<LocalShell>>);

impl LocalShellScope {
    fn new(arc: &Arc<Mutex<LocalShell>>) -> LocalShellScope {
        let mut lock = PROCESS_MANAGER.lock().unwrap();
        let id = thread::current().id();
        lock.add_local_shell(&id, arc);

        LocalShellScope(id, arc.clone())
    }
}

impl Default for LocalShellScope {
    fn default() -> LocalShellScope {
        let id = thread::current().id();
        LocalShellScope(id, Arc::new(Mutex::new(LocalShell::new())))
    }
}

impl Drop for LocalShellScope {
    fn drop(&mut self) {
        let mut lock = PROCESS_MANAGER.lock().unwrap();
        lock.remove_local_shell(&self.0);
    }
}

pub struct ShellHandle<T> {
    join_handle: JoinHandle<T>,
    shell: Arc<Mutex<LocalShell>>
}

impl <T> ShellHandle<T> {
    pub fn signal(&self, signal: c_int) {
        let mut lock = self.shell.lock().unwrap();
        lock.signal(signal);
    }

    pub fn terminate(self) -> Result<T, Box<Any + Send + 'static>> {
        self.signal(libc::SIGTERM);
        self.join_handle.join()
    }
}

pub fn spawn<F, T>(f: F) -> ShellHandle<T> where
        F: FnOnce() -> T, F: Send + 'static, T: Send + 'static {
    let arc = Arc::new(Mutex::new(LocalShell::new()));
    let arc_clone = arc.clone();
    let join_handle = thread::spawn(move || -> T {
        LOCAL_SHELL_SCOPE.with(|shell| {
            let mut shell = shell.borrow_mut();
            if shell.is_some() {
                panic!("Shell has already registered");
            }
            *shell = Some(LocalShellScope::new(&arc_clone));
        });
        f()
    });
    ShellHandle {
        join_handle: join_handle,
        shell: arc
    }
}

pub fn current_shell() -> Arc<Mutex<LocalShell>> {
    LOCAL_SHELL_SCOPE.with(|shell| {
        shell.borrow_mut()
            .get_or_insert(LocalShellScope::default()).1.clone()
    })
}

thread_local! {
    static LOCAL_SHELL_SCOPE: RefCell<Option<LocalShellScope>> =
        RefCell::new(None);
}

<commit_msg>Implement join for ShellHandle<commit_after>use job_handle::ChildProcess;
use libc::c_int;
use libc;
use process_manager::PROCESS_MANAGER;
use result::ShellResultExt;
use std::any::Any;
use std::cell::RefCell;
use std::ops::Deref;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::RwLock;
use std::thread::JoinHandle;
use std::thread::ThreadId;
use std::thread;

/// Thread local shell.
pub struct LocalShell {
    thread_id: ThreadId,
    processes: Vec<Arc<RwLock<ChildProcess>>>,
    signaled: bool
}

impl LocalShell {
    fn new() -> LocalShell {
        LocalShell {
            thread_id: thread::current().id(),
            processes: Vec::new(),
            signaled: false
        }
    }

    pub fn add_process(&mut self, process: &Arc<RwLock<ChildProcess>>) {
        self.processes.push(process.clone());
    }

    pub fn remove_process(&mut self, process: &Arc<RwLock<ChildProcess>>) {
        self.processes.retain(|p| !Arc::ptr_eq(p, process));
    }

    pub fn signal(&mut self, signal: c_int) {
        self.signaled = true;
        for process in &self.processes {
            let lock = process.read().unwrap();
            lock.signal(signal).print_error();
        }
    }

    pub fn signaled(&self) -> bool {
        self.signaled
    }

    pub fn thread_id(&self) -> &ThreadId {
        &self.thread_id
    }
}

struct LocalShellScope(ThreadId, Arc<Mutex<LocalShell>>);

impl LocalShellScope {
    fn new(arc: &Arc<Mutex<LocalShell>>) -> LocalShellScope {
        let mut lock = PROCESS_MANAGER.lock().unwrap();
        let id = thread::current().id();
        lock.add_local_shell(&id, arc);

        LocalShellScope(id, arc.clone())
    }
}

impl Default for LocalShellScope {
    fn default() -> LocalShellScope {
        let id = thread::current().id();
        LocalShellScope(id, Arc::new(Mutex::new(LocalShell::new())))
    }
}

impl Drop for LocalShellScope {
    fn drop(&mut self) {
        let mut lock = PROCESS_MANAGER.lock().unwrap();
        lock.remove_local_shell(&self.0);
    }
}

pub struct ShellHandle<T> {
    join_handle: JoinHandle<T>,
    shell: Arc<Mutex<LocalShell>>
}

impl <T> ShellHandle<T> {
    pub fn signal(&self, signal: c_int) {
        let mut lock = self.shell.lock().unwrap();
        lock.signal(signal);
    }

    pub fn terminate(self) -> Result<T, Box<Any + Send + 'static>> {
        self.signal(libc::SIGTERM);
        self.join_handle.join()
    }

    pub fn join(self) -> Result<T, Box<Any + Send + 'static>> {
        self.join_handle.join()
    }
}

impl <T> Deref for ShellHandle<T> {
    type Target = JoinHandle<T>;
    fn deref(&self) -> &Self::Target {
        &self.join_handle
    }
}

pub fn spawn<F, T>(f: F) -> ShellHandle<T> where
        F: FnOnce() -> T, F: Send + 'static, T: Send + 'static {
    let arc = Arc::new(Mutex::new(LocalShell::new()));
    let arc_clone = arc.clone();
    let join_handle = thread::spawn(move || -> T {
        LOCAL_SHELL_SCOPE.with(|shell| {
            let mut shell = shell.borrow_mut();
            if shell.is_some() {
                panic!("Shell has already registered");
            }
            *shell = Some(LocalShellScope::new(&arc_clone));
        });
        f()
    });
    ShellHandle {
        join_handle: join_handle,
        shell: arc
    }
}

pub fn current_shell() -> Arc<Mutex<LocalShell>> {
    LOCAL_SHELL_SCOPE.with(|shell| {
        shell.borrow_mut()
            .get_or_insert(LocalShellScope::default()).1.clone()
    })
}

thread_local! {
    static LOCAL_SHELL_SCOPE: RefCell<Option<LocalShellScope>> =
        RefCell::new(None);
}

<|endoftext|>


<commit_before>extern crate num;
extern crate rand;
extern crate bignum = "bignum#0.1.1-pre";

use bignum::{BigUint, ToBigUint,RandBigInt};
use num::{div_mod_floor, Integer};
use std::num::{Zero, One, pow, Float};
use rand::{task_rng, Rng};
use rand::distributions::range::SampleRange;
use std::vec::Vec;
use std::iter;
use std::clone::Clone;
use std::ops::{Shr, BitAnd};

#[inline]
fn big<T: ToBigUint>(i : T) -> BigUint {
    i.to_biguint().unwrap()
}

// Source: https://gist.github.com/jsanders/8739134#file-generate_big_primes-rs-L35
fn mod_exp<T: Integer + Clone + Shr<uint,T> + BitAnd<T,T>>
    (base: &T, exponent: &T, modulus: &T) -> T {

    let (zero, one): (T, T) = (Zero::zero(), One::one());
    let mut result = one.clone();
    let mut baseAcc = base.clone();
    let mut exponentAcc = exponent.clone();

    while exponentAcc > zero {
        // Accumulate current base if current exponent bit is 1
        if (exponentAcc & one) == one {
            result = result.mul(&baseAcc);
            result = result.rem(modulus);
        }
        // Get next base by squaring
        baseAcc = baseAcc * baseAcc;
        baseAcc = baseAcc % *modulus;

        // Get next bit of exponent
        exponentAcc = exponentAcc.shr(&1);
    }
    result

}

// Miller-Rabin, probabilistic
fn try_composite<T : Integer + Clone + Shr<uint,T> + BitAnd<T,T> + ToPrimitive>
    (a : &T, d : &T, n : &T, s : &T) -> bool {

    let one : T = One::one();
    if mod_exp(a, d, n) == one {
        return false
    }
    let mut i = 0u;
    loop {
        if i == s.to_uint().unwrap() {
            break
        }
        if mod_exp(a, &(pow(one+one, i) * *d), n) == *n-one {
            return false
        }
        i += 1;
    }
    true

}
fn is_prime<T : Integer + Clone + Shr<uint,T> + BitAnd<T,T> + ToPrimitive + SampleRange>
    (n : &T) -> bool {

    let (zero, one): (T, T) = (Zero::zero(), One::one());
    let two = one + one;

    assert!(n >= &two);

    let mut s = zero.clone();
    let mut d = *n-one;
    let num_trials = 5;
    if n == &two {
        return true
    } else if n == &one || n.is_even() {
        return false
    } else {
        loop {
            let (quotient, remainder) = div_mod_floor(d.clone(), two.clone());
            if remainder == one { break }
            s = s + one;
            d = quotient;
        }
    }

    let mut rng = task_rng();
    for _ in range(0, num_trials) {
        let a = rng.gen_range(one + one, n.clone());
        if try_composite(&a, &d, n, &s) {
            return false
        }
    }
    true

}

fn is_valid_pow(prime : uint) -> bool {
    for &offset in [0u, 4u, 6u, 10u, 12u, 16u].iter() {
        if !fermat_is_prime(prime + offset) {
            return false;
        }
    }
    true
}

fn fermat_is_prime(n : uint) -> bool {
    if n == 2 {
        true
    } else if n.is_even() {
        false
    } else if mod_exp(&2, &(n-1), &n) == 1 {
        true
    } else {
        false
    }
}

fn gen_prime(max_val : uint, verbose : bool) -> uint {
    //http://rosettacode.org/wiki/Sieve_of_Eratosthenes#Rust
    fn simple_sieve(limit: uint) -> ~[uint] {
        #[inline]
        fn int_sqrt(n: uint) -> uint { (n as f64).sqrt() as uint }

        if limit < 2 {
            return ~[];
        }

        let mut primes = Vec::from_elem(limit + 1, true);

        for prime in iter::range_inclusive(2, int_sqrt(limit) + 1) {
            if *primes.get(prime) {
                for multiple in iter::range_step(prime * prime, limit + 1, prime) {
                    *primes.get_mut(multiple) = false
                }
            }
        }
        iter::range_inclusive(2, limit).filter(|&n| *primes.get(n)).collect()
    }

    fn candidate_killed_by (candidate : uint, prime : uint) -> bool {
        let offsets : Vec<uint> = Vec::from_slice([0u, 4u, 6u, 10u, 12u, 16u]);
        for offset in offsets.iter() {
            if (candidate + *offset) % prime == 0 {
                return true;
            }
        }
        false
    }

    fn add_next_prime (max_val : uint, offsets : Vec<uint>, prime : uint, primorial : uint) -> Vec<uint> {
        let mut base : uint = 0u;
        let mut new_offsets : Vec<uint> = Vec::new();
        for _ in range(0u, prime) {
            if base > max_val {
                break
            }
            for o in offsets.iter() {
                    let val = base + *o;
                    if val > max_val {
                        break
                    }
                    if !candidate_killed_by(val, prime) {
                        new_offsets.push(val);
                    }
            }
            base = base + primorial;
        }
        return new_offsets
    }
    let primorial_max : uint = 29u;
    let primorial_start : uint = 7u;
    let mut primorial : uint = 210u;  // 2*3*5*7
    let mut offsets : Vec<uint> = Vec::from_slice([97u]);
    for i in simple_sieve(primorial_max).iter() {
        if *i <= primorial_start {
            continue
        }
        offsets = add_next_prime(max_val, offsets, *i, primorial);
        primorial = primorial * *i;
    }
    let mut count : uint = 0;
    for &o in offsets.iter() {
        if is_valid_pow(o) {
            count = count + 1;
            if verbose { println!("prime: {}", o) }
        }
    }
    count
}

#[cfg(test)]
mod test_primes {
    use super::gen_prime;

    #[test]
    fn gen_prime_is_correct() {
        let v = gen_prime(1000000000, true);
        println!("{}", v);
        assert!(v == 81u);
    }
}

fn main() {
    println!("{}", gen_prime(1000000000, true));
}<commit_msg>Fully working BigUint POW<commit_after>extern crate num;
extern crate rand;
extern crate bignum = "bignum#0.1.1-pre";

use bignum::{BigUint, ToBigUint, RandBigInt};
use num::{div_mod_floor, Integer};
use std::num::{Zero, One, pow, Float, FromPrimitive};
use rand::{task_rng, Rng};
use rand::distributions::range::SampleRange;
use std::vec::Vec;
use std::iter;
use std::clone::Clone;
use std::ops::{Shr, BitAnd};

static pow_offsets : &'static[uint] = &[0u, 4u, 6u, 10u, 12u, 16u];

#[inline]
fn big<T: ToBigUint>(i : T) -> BigUint {
    i.to_biguint().unwrap()
}

// Source: https://gist.github.com/jsanders/8739134#file-generate_big_primes-rs-L35
fn mod_exp<T: Integer + Clone + Shr<uint,T> + BitAnd<T,T>>
    (base: &T, exponent: &T, modulus: &T) -> T {

    let (zero, one): (T, T) = (Zero::zero(), One::one());
    let mut result = one.clone();
    let mut baseAcc = base.clone();
    let mut exponentAcc = exponent.clone();

    while exponentAcc > zero {
        // Accumulate current base if current exponent bit is 1
        if (exponentAcc & one) == one {
            result = result.mul(&baseAcc);
            result = result.rem(modulus);
        }
        // Get next base by squaring
        baseAcc = baseAcc * baseAcc;
        baseAcc = baseAcc % *modulus;

        // Get next bit of exponent
        exponentAcc = exponentAcc.shr(&1);
    }
    result

}

// Miller-Rabin, probabilistic
fn try_composite(a : &BigUint, d : &BigUint, n : &BigUint, s : &BigUint) -> bool {

    let one : BigUint = One::one();
    if mod_exp(a, d, n) == one {
        return false
    }
    let mut i = 0u;
    loop {
        if big(i) == *s {
            break
        }
        if mod_exp(a, &(pow(one+one, i) * *d), n) == *n-one {
            return false
        }
        i += 1;
    }
    true

}
fn is_prime(n : &BigUint) -> bool {

    let (zero, one): (BigUint, BigUint) = (Zero::zero(), One::one());
    let two = one + one;

    assert!(n >= &two);

    let mut s = zero.clone();
    let mut d = *n-one;
    let num_trials = 5;
    if n == &two {
        return true
    } else if n == &one || n.is_even() {
        return false
    } else {
        loop {
            let (quotient, remainder) = div_mod_floor(d.clone(), two.clone());
            if remainder == one { break }
            s = s + one;
            d = quotient;
        }
    }

    let mut rng = task_rng();
    for _ in range(0, num_trials) {
        let a = rng.gen_biguint_range(&two, n);
        if try_composite(&a, &d, n, &s) {
            return false
        }
    }
    true

}

fn fermat_is_prime(n : &BigUint) -> bool {
    let (zero, one): (BigUint, BigUint) = (Zero::zero(), One::one());
    let two = one + one;

    if n == &two {
        true
    } else if n.is_even() {
        false
    } else if mod_exp(&two, &(n-one), n) == one {
        true
    } else {
        false
    }
}

fn is_valid_pow(prime : &BigUint) -> bool {
    for offset in [big(0u), big(4u), big(6u), big(10u), big(12u), big(16u)].iter() {
        if !is_prime(&(*prime + *offset)) {
            return false;
        }
    }
    true
}

// Source: http://rosettacode.org/wiki/Sieve_of_Eratosthenes#Rust
fn simple_sieve(limit: uint) -> ~[uint] {

    #[inline]
    fn int_sqrt(n: uint) -> uint { (n as f64).sqrt() as uint }

    if limit < 2 {
        return ~[];
    }

    let mut primes = Vec::from_elem(limit + 1, true);

    for prime in iter::range_inclusive(2, int_sqrt(limit) + 1) {
        if *primes.get(prime) {
            for multiple in iter::range_step(prime * prime, limit + 1, prime) {
                *primes.get_mut(multiple) = false
            }
        }
    }
    iter::range_inclusive(2, limit).filter(|&n| *primes.get(n)).collect()

}

fn gen_prime(max_val : &BigUint, verbose : bool) -> uint {

    fn candidate_killed_by(candidate : &BigUint, prime : &BigUint) -> bool {
        let zero : BigUint = Zero::zero();
        for &offset in pow_offsets.iter() {
            let o : BigUint = FromPrimitive::from_uint(offset).unwrap();
            if (*candidate + o) % *prime == zero {
                return true;
            }
        }
        false
    }

    fn add_next_prime (max_val : &BigUint, offsets : Vec<BigUint>, prime : BigUint, primorial : BigUint) -> Vec<BigUint> {
        let (zero, one): (BigUint, BigUint) = (Zero::zero(), One::one());
        let mut base : BigUint = Zero::zero();
        let mut new_offsets : Vec<BigUint> = Vec::new();
        let mut counter = zero.clone();
        while counter < prime {
            if base > *max_val {
                break
            }
            for o in offsets.iter() {
                    let val = base + *o;
                    if val > *max_val {
                        break
                    }
                    if !candidate_killed_by(&val, &prime) {
                        new_offsets.push(val);
                    }
            }
            base = base + primorial;
            counter = counter + one;
        }
        return new_offsets
    }
    let primorial_max : uint = 29u;
    let primorial_start : uint = 7u;
    let mut primorial : uint = 210u;  // 2*3*5*7
    let mut offsets : Vec<BigUint> = Vec::from_slice([big(97u)]);
    for i in simple_sieve(primorial_max).iter() {
        if *i <= primorial_start {
            continue
        }
        offsets = add_next_prime(max_val, offsets, big(*i), big(primorial));
        primorial = primorial * *i;
    }
    let mut count : uint = 0;
    for o in offsets.iter() {
        if is_valid_pow(o) {
            count = count + 1;
            if verbose { println!("prime: {}", o) }
        }
    }
    count
}

#[cfg(test)]
mod test_primes {
    use super::gen_prime;

    #[test]
    fn gen_prime_is_correct() {
        let v = gen_prime(&big(1000000000), true);
        println!("{}", v);
        assert!(v == 81u);
    }
}

fn main() {
    println!("{}", gen_prime(&big(1000000000), false));
}<|endoftext|>


<commit_before>/// Components are types that associate an entity with a set of properties
/// for each type a list of all its instances is kept so they can be iterated
/// over efficiently by systems.
///
/// Whenever a new component is created it is added to a list.
///
///
///
/// component! { physics, body: physics.RigidBody, physics_id: physics.ID }
///

use std::collections::{ HashSet, HashMap };
use std::any::{ Any, TypeId };
use shared_mutex::{ SharedMutex };

pub mod example;

#[derive(Clone)]
pub struct Component {
	pub name: TypeId,
	pub get_component_list: fn () -> Box<Any>
}

lazy_static! {
	pub static ref COMPONENTS: SharedMutex<HashMap<TypeId, Component>> = SharedMutex::new(HashMap::new());
}

pub fn register(component : Component) {
	let mut components = COMPONENTS.write().expect("COMPONENTS lock corrupted");
	components.insert(component.name, component);
}

pub fn get_component_mutex<T>(id : TypeId) -> T {

}

#[macro_export]
macro_rules! component {
	( $component_name:ident , $( $name:ident : $field:ty ),* ) => (
		pub mod $component_name {
			use shared_mutex::{ SharedMutex, SharedMutexWriteGuard };
			use entity_rust::entities::{ ComponentList, EntityID };

			#[derive(Default)]
			pub struct Component {
				pub $($name : $field),*
			}

			lazy_static! {
				pub static ref LIST: SharedMutex<ComponentList<Component>> = SharedMutex::new(Vec::new());
			}

			pub fn create(mut list: SharedMutexWriteGuard<ComponentList<Component>>, entity: EntityID, $($name : $field),*) {
				let c = Component { $($name : $name),* };
				list.push((entity,c));
			}
		}
	)
}
<commit_msg>try and implement the lock getter<commit_after>/// Components are types that associate an entity with a set of properties
/// for each type a list of all its instances is kept so they can be iterated
/// over efficiently by systems.
///
/// Whenever a new component is created it is added to a list.
///
/// component! { physics, body: physics.RigidBody, physics_id: physics.ID }
///

use std::collections::{ HashSet, HashMap };
use std::any::{ Any, TypeId };
use shared_mutex::{ SharedMutex, SharedMutexReadGuard, SharedMutexWriteGuard };

pub mod example;

#[derive(Clone)]
pub struct Component {
	pub name: TypeId,
	pub get_component_list: fn () -> Box<Any + 'static>
}

lazy_static! {
	pub static ref COMPONENTS: SharedMutex<HashMap<TypeId, Component>> = SharedMutex::new(HashMap::new());
}

pub fn register(component : Component) {
	let mut components = COMPONENTS.write().expect("COMPONENTS lock corrupted");
	components.insert(component.name, component);
}

pub fn get_components_lock<'mutex, T>(id : TypeId) -> SharedMutexReadGuard<'mutex, T> where T : 'static {
	let components = COMPONENTS.read().expect("COMPONENTS lock corrupted");
	let component = components.get(&id).expect("Unknown component type requested");
	let component_list_any = (component.get_component_list)();
	let component_list : &'static SharedMutex<T> = *component_list_any.downcast_ref().expect("Invalid type for components list");
	return component_list.read().expect("Component LIST lock corrupted");
}

#[macro_export]
macro_rules! component {
	( $component_name:ident , $( $name:ident : $field:ty ),* ) => (
		pub mod $component_name {
			use shared_mutex::{ SharedMutex, SharedMutexWriteGuard };
			use entity_rust::entities::{ ComponentList, EntityID };

			#[derive(Default)]
			pub struct Component {
				pub $($name : $field),*
			}

			lazy_static! {
				pub static ref LIST: SharedMutex<ComponentList<Component>> = SharedMutex::new(Vec::new());
			}

			pub fn create(mut list: SharedMutexWriteGuard<ComponentList<Component>>, entity: EntityID, $($name : $field),*) {
				let c = Component { $($name : $name),* };
				list.push((entity,c));
			}
		}
	)
}
<|endoftext|>


<commit_before>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

/*!
Sampling from random distributions.

This is a generalization of `Rand` to allow parameters to control the
exact properties of the generated values, e.g. the mean and standard
deviation of a normal distribution. The `Sample` trait is the most
general, and allows for generating values that change some state
internally. The `IndependentSample` trait is for generating values
that do not need to record state.

*/

#![experimental]

use core::prelude::*;
use core::num::Int;

use {Rng, Rand};

pub use self::range::Range;
pub use self::gamma::{Gamma, ChiSquared, FisherF, StudentT};
pub use self::normal::{Normal, LogNormal};
pub use self::exponential::Exp;

pub mod range;
pub mod gamma;
pub mod normal;
pub mod exponential;

/// Types that can be used to create a random instance of `Support`.
pub trait Sample<Support> {
    /// Generate a random value of `Support`, using `rng` as the
    /// source of randomness.
    fn sample<R: Rng>(&mut self, rng: &mut R) -> Support;
}

/// `Sample`s that do not require keeping track of state.
///
/// Since no state is recorded, each sample is (statistically)
/// independent of all others, assuming the `Rng` used has this
/// property.
// FIXME maybe having this separate is overkill (the only reason is to
// take &self rather than &mut self)? or maybe this should be the
// trait called `Sample` and the other should be `DependentSample`.
pub trait IndependentSample<Support>: Sample<Support> {
    /// Generate a random value.
    fn ind_sample<R: Rng>(&self, &mut R) -> Support;
}

/// A wrapper for generating types that implement `Rand` via the
/// `Sample` & `IndependentSample` traits.
pub struct RandSample<Sup>;

impl<Sup: Rand> Sample<Sup> for RandSample<Sup> {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> Sup { self.ind_sample(rng) }
}

impl<Sup: Rand> IndependentSample<Sup> for RandSample<Sup> {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> Sup {
        rng.gen()
    }
}

/// A value with a particular weight for use with `WeightedChoice`.
pub struct Weighted<T> {
    /// The numerical weight of this item
    pub weight: uint,
    /// The actual item which is being weighted
    pub item: T,
}

/// A distribution that selects from a finite collection of weighted items.
///
/// Each item has an associated weight that influences how likely it
/// is to be chosen: higher weight is more likely.
///
/// The `Clone` restriction is a limitation of the `Sample` and
/// `IndependentSample` traits. Note that `&T` is (cheaply) `Clone` for
/// all `T`, as is `uint`, so one can store references or indices into
/// another vector.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{Weighted, WeightedChoice, IndependentSample};
///
/// let mut items = vec!(Weighted { weight: 2, item: 'a' },
///                      Weighted { weight: 4, item: 'b' },
///                      Weighted { weight: 1, item: 'c' });
/// let wc = WeightedChoice::new(items.as_mut_slice());
/// let mut rng = rand::task_rng();
/// for _ in range(0u, 16) {
///      // on average prints 'a' 4 times, 'b' 8 and 'c' twice.
///      println!("{}", wc.ind_sample(&mut rng));
/// }
/// ```
pub struct WeightedChoice<'a, T:'a> {
    items: &'a mut [Weighted<T>],
    weight_range: Range<uint>
}

impl<'a, T: Clone> WeightedChoice<'a, T> {
    /// Create a new `WeightedChoice`.
    ///
    /// Fails if:
    /// - `v` is empty
    /// - the total weight is 0
    /// - the total weight is larger than a `uint` can contain.
    pub fn new<'a>(items: &'a mut [Weighted<T>]) -> WeightedChoice<'a, T> {
        // strictly speaking, this is subsumed by the total weight == 0 case
        assert!(!items.is_empty(), "WeightedChoice::new called with no items");

        let mut running_total = 0u;

        // we convert the list from individual weights to cumulative
        // weights so we can binary search. This *could* drop elements
        // with weight == 0 as an optimisation.
        for item in items.iter_mut() {
            running_total = match running_total.checked_add(item.weight) {
                Some(n) => n,
                None => panic!("WeightedChoice::new called with a total weight \
                               larger than a uint can contain")
            };

            item.weight = running_total;
        }
        assert!(running_total != 0, "WeightedChoice::new called with a total weight of 0");

        WeightedChoice {
            items: items,
            // we're likely to be generating numbers in this range
            // relatively often, so might as well cache it
            weight_range: Range::new(0, running_total)
        }
    }
}

impl<'a, T: Clone> Sample<T> for WeightedChoice<'a, T> {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> T { self.ind_sample(rng) }
}

impl<'a, T: Clone> IndependentSample<T> for WeightedChoice<'a, T> {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> T {
        // we want to find the first element that has cumulative
        // weight > sample_weight, which we do by binary since the
        // cumulative weights of self.items are sorted.

        // choose a weight in [0, total_weight)
        let sample_weight = self.weight_range.ind_sample(rng);

        // short circuit when it's the first item
        if sample_weight < self.items[0].weight {
            return self.items[0].item.clone();
        }

        let mut idx = 0;
        let mut modifier = self.items.len();

        // now we know that every possibility has an element to the
        // left, so we can just search for the last element that has
        // cumulative weight <= sample_weight, then the next one will
        // be "it". (Note that this greatest element will never be the
        // last element of the vector, since sample_weight is chosen
        // in [0, total_weight) and the cumulative weight of the last
        // one is exactly the total weight.)
        while modifier > 1 {
            let i = idx + modifier / 2;
            if self.items[i].weight <= sample_weight {
                // we're small, so look to the right, but allow this
                // exact element still.
                idx = i;
                // we need the `/ 2` to round up otherwise we'll drop
                // the trailing elements when `modifier` is odd.
                modifier += 1;
            } else {
                // otherwise we're too big, so go left. (i.e. do
                // nothing)
            }
            modifier /= 2;
        }
        return self.items[idx + 1].item.clone();
    }
}

mod ziggurat_tables;

/// Sample a random number using the Ziggurat method (specifically the
/// ZIGNOR variant from Doornik 2005). Most of the arguments are
/// directly from the paper:
///
/// * `rng`: source of randomness
/// * `symmetric`: whether this is a symmetric distribution, or one-sided with P(x < 0) = 0.
/// * `X`: the $x_i$ abscissae.
/// * `F`: precomputed values of the PDF at the $x_i$, (i.e. $f(x_i)$)
/// * `F_DIFF`: precomputed values of $f(x_i) - f(x_{i+1})$
/// * `pdf`: the probability density function
/// * `zero_case`: manual sampling from the tail when we chose the
///    bottom box (i.e. i == 0)

// the perf improvement (25-50%) is definitely worth the extra code
// size from force-inlining.
#[inline(always)]
fn ziggurat<R:Rng>(
            rng: &mut R,
            symmetric: bool,
            x_tab: ziggurat_tables::ZigTable,
            f_tab: ziggurat_tables::ZigTable,
            pdf: |f64|: 'static -> f64,
            zero_case: |&mut R, f64|: 'static -> f64)
            -> f64 {
    static SCALE: f64 = (1u64 << 53) as f64;
    loop {
        // reimplement the f64 generation as an optimisation suggested
        // by the Doornik paper: we have a lot of precision-space
        // (i.e. there are 11 bits of the 64 of a u64 to use after
        // creating a f64), so we might as well reuse some to save
        // generating a whole extra random number. (Seems to be 15%
        // faster.)
        //
        // This unfortunately misses out on the benefits of direct
        // floating point generation if an RNG like dSMFT is
        // used. (That is, such RNGs create floats directly, highly
        // efficiently and overload next_f32/f64, so by not calling it
        // this may be slower than it would be otherwise.)
        // FIXME: investigate/optimise for the above.
        let bits: u64 = rng.gen();
        let i = (bits & 0xff) as uint;
        let f = (bits >> 11) as f64 / SCALE;

        // u is either U(-1, 1) or U(0, 1) depending on if this is a
        // symmetric distribution or not.
        let u = if symmetric {2.0 * f - 1.0} else {f};
        let x = u * x_tab[i];

        let test_x = if symmetric { x.abs() } else {x};

        // algebraically equivalent to |u| < x_tab[i+1]/x_tab[i] (or u < x_tab[i+1]/x_tab[i])
        if test_x < x_tab[i + 1] {
            return x;
        }
        if i == 0 {
            return zero_case(rng, u);
        }
        // algebraically equivalent to f1 + DRanU()*(f0 - f1) < 1
        if f_tab[i + 1] + (f_tab[i] - f_tab[i + 1]) * rng.gen() < pdf(x) {
            return x;
        }
    }
}

#[cfg(test)]
mod tests {
    use std::prelude::*;

    use {Rng, Rand};
    use super::{RandSample, WeightedChoice, Weighted, Sample, IndependentSample};

    #[deriving(PartialEq, Show)]
    struct ConstRand(uint);
    impl Rand for ConstRand {
        fn rand<R: Rng>(_: &mut R) -> ConstRand {
            ConstRand(0)
        }
    }

    // 0, 1, 2, 3, ...
    struct CountingRng { i: u32 }
    impl Rng for CountingRng {
        fn next_u32(&mut self) -> u32 {
            self.i += 1;
            self.i - 1
        }
        fn next_u64(&mut self) -> u64 {
            self.next_u32() as u64
        }
    }

    #[test]
    fn test_rand_sample() {
        let mut rand_sample = RandSample::<ConstRand>;

        assert_eq!(rand_sample.sample(&mut ::test::rng()), ConstRand(0));
        assert_eq!(rand_sample.ind_sample(&mut ::test::rng()), ConstRand(0));
    }
    #[test]
    fn test_weighted_choice() {
        // this makes assumptions about the internal implementation of
        // WeightedChoice, specifically: it doesn't reorder the items,
        // it doesn't do weird things to the RNG (so 0 maps to 0, 1 to
        // 1, internally; modulo a modulo operation).

        macro_rules! t (
            ($items:expr, $expected:expr) => {{
                let mut items = $items;
                let wc = WeightedChoice::new(items.as_mut_slice());
                let expected = $expected;

                let mut rng = CountingRng { i: 0 };

                for &val in expected.iter() {
                    assert_eq!(wc.ind_sample(&mut rng), val)
                }
            }}
        );

        t!(vec!(Weighted { weight: 1, item: 10i}), [10]);

        // skip some
        t!(vec!(Weighted { weight: 0, item: 20i},
                Weighted { weight: 2, item: 21i},
                Weighted { weight: 0, item: 22i},
                Weighted { weight: 1, item: 23i}),
           [21,21, 23]);

        // different weights
        t!(vec!(Weighted { weight: 4, item: 30i},
                Weighted { weight: 3, item: 31i}),
           [30,30,30,30, 31,31,31]);

        // check that we're binary searching
        // correctly with some vectors of odd
        // length.
        t!(vec!(Weighted { weight: 1, item: 40i},
                Weighted { weight: 1, item: 41i},
                Weighted { weight: 1, item: 42i},
                Weighted { weight: 1, item: 43i},
                Weighted { weight: 1, item: 44i}),
           [40, 41, 42, 43, 44]);
        t!(vec!(Weighted { weight: 1, item: 50i},
                Weighted { weight: 1, item: 51i},
                Weighted { weight: 1, item: 52i},
                Weighted { weight: 1, item: 53i},
                Weighted { weight: 1, item: 54i},
                Weighted { weight: 1, item: 55i},
                Weighted { weight: 1, item: 56i}),
           [50, 51, 52, 53, 54, 55, 56]);
    }

    #[test] #[should_fail]
    fn test_weighted_choice_no_items() {
        WeightedChoice::<int>::new([]);
    }
    #[test] #[should_fail]
    fn test_weighted_choice_zero_weight() {
        WeightedChoice::new(&mut [Weighted { weight: 0, item: 0i},
                                  Weighted { weight: 0, item: 1i}]);
    }
    #[test] #[should_fail]
    fn test_weighted_choice_weight_overflows() {
        let x = (-1) as uint / 2; // x + x + 2 is the overflow
        WeightedChoice::new(&mut [Weighted { weight: x, item: 0i },
                                  Weighted { weight: 1, item: 1i },
                                  Weighted { weight: x, item: 2i },
                                  Weighted { weight: 1, item: 3i }]);
    }
}
<commit_msg>Remove Signed trait and add SignedInt trait<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

/*!
Sampling from random distributions.

This is a generalization of `Rand` to allow parameters to control the
exact properties of the generated values, e.g. the mean and standard
deviation of a normal distribution. The `Sample` trait is the most
general, and allows for generating values that change some state
internally. The `IndependentSample` trait is for generating values
that do not need to record state.

*/

#![experimental]

use core::prelude::*;
use core::num::{Float, Int};

use {Rng, Rand};

pub use self::range::Range;
pub use self::gamma::{Gamma, ChiSquared, FisherF, StudentT};
pub use self::normal::{Normal, LogNormal};
pub use self::exponential::Exp;

pub mod range;
pub mod gamma;
pub mod normal;
pub mod exponential;

/// Types that can be used to create a random instance of `Support`.
pub trait Sample<Support> {
    /// Generate a random value of `Support`, using `rng` as the
    /// source of randomness.
    fn sample<R: Rng>(&mut self, rng: &mut R) -> Support;
}

/// `Sample`s that do not require keeping track of state.
///
/// Since no state is recorded, each sample is (statistically)
/// independent of all others, assuming the `Rng` used has this
/// property.
// FIXME maybe having this separate is overkill (the only reason is to
// take &self rather than &mut self)? or maybe this should be the
// trait called `Sample` and the other should be `DependentSample`.
pub trait IndependentSample<Support>: Sample<Support> {
    /// Generate a random value.
    fn ind_sample<R: Rng>(&self, &mut R) -> Support;
}

/// A wrapper for generating types that implement `Rand` via the
/// `Sample` & `IndependentSample` traits.
pub struct RandSample<Sup>;

impl<Sup: Rand> Sample<Sup> for RandSample<Sup> {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> Sup { self.ind_sample(rng) }
}

impl<Sup: Rand> IndependentSample<Sup> for RandSample<Sup> {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> Sup {
        rng.gen()
    }
}

/// A value with a particular weight for use with `WeightedChoice`.
pub struct Weighted<T> {
    /// The numerical weight of this item
    pub weight: uint,
    /// The actual item which is being weighted
    pub item: T,
}

/// A distribution that selects from a finite collection of weighted items.
///
/// Each item has an associated weight that influences how likely it
/// is to be chosen: higher weight is more likely.
///
/// The `Clone` restriction is a limitation of the `Sample` and
/// `IndependentSample` traits. Note that `&T` is (cheaply) `Clone` for
/// all `T`, as is `uint`, so one can store references or indices into
/// another vector.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{Weighted, WeightedChoice, IndependentSample};
///
/// let mut items = vec!(Weighted { weight: 2, item: 'a' },
///                      Weighted { weight: 4, item: 'b' },
///                      Weighted { weight: 1, item: 'c' });
/// let wc = WeightedChoice::new(items.as_mut_slice());
/// let mut rng = rand::task_rng();
/// for _ in range(0u, 16) {
///      // on average prints 'a' 4 times, 'b' 8 and 'c' twice.
///      println!("{}", wc.ind_sample(&mut rng));
/// }
/// ```
pub struct WeightedChoice<'a, T:'a> {
    items: &'a mut [Weighted<T>],
    weight_range: Range<uint>
}

impl<'a, T: Clone> WeightedChoice<'a, T> {
    /// Create a new `WeightedChoice`.
    ///
    /// Fails if:
    /// - `v` is empty
    /// - the total weight is 0
    /// - the total weight is larger than a `uint` can contain.
    pub fn new<'a>(items: &'a mut [Weighted<T>]) -> WeightedChoice<'a, T> {
        // strictly speaking, this is subsumed by the total weight == 0 case
        assert!(!items.is_empty(), "WeightedChoice::new called with no items");

        let mut running_total = 0u;

        // we convert the list from individual weights to cumulative
        // weights so we can binary search. This *could* drop elements
        // with weight == 0 as an optimisation.
        for item in items.iter_mut() {
            running_total = match running_total.checked_add(item.weight) {
                Some(n) => n,
                None => panic!("WeightedChoice::new called with a total weight \
                               larger than a uint can contain")
            };

            item.weight = running_total;
        }
        assert!(running_total != 0, "WeightedChoice::new called with a total weight of 0");

        WeightedChoice {
            items: items,
            // we're likely to be generating numbers in this range
            // relatively often, so might as well cache it
            weight_range: Range::new(0, running_total)
        }
    }
}

impl<'a, T: Clone> Sample<T> for WeightedChoice<'a, T> {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> T { self.ind_sample(rng) }
}

impl<'a, T: Clone> IndependentSample<T> for WeightedChoice<'a, T> {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> T {
        // we want to find the first element that has cumulative
        // weight > sample_weight, which we do by binary since the
        // cumulative weights of self.items are sorted.

        // choose a weight in [0, total_weight)
        let sample_weight = self.weight_range.ind_sample(rng);

        // short circuit when it's the first item
        if sample_weight < self.items[0].weight {
            return self.items[0].item.clone();
        }

        let mut idx = 0;
        let mut modifier = self.items.len();

        // now we know that every possibility has an element to the
        // left, so we can just search for the last element that has
        // cumulative weight <= sample_weight, then the next one will
        // be "it". (Note that this greatest element will never be the
        // last element of the vector, since sample_weight is chosen
        // in [0, total_weight) and the cumulative weight of the last
        // one is exactly the total weight.)
        while modifier > 1 {
            let i = idx + modifier / 2;
            if self.items[i].weight <= sample_weight {
                // we're small, so look to the right, but allow this
                // exact element still.
                idx = i;
                // we need the `/ 2` to round up otherwise we'll drop
                // the trailing elements when `modifier` is odd.
                modifier += 1;
            } else {
                // otherwise we're too big, so go left. (i.e. do
                // nothing)
            }
            modifier /= 2;
        }
        return self.items[idx + 1].item.clone();
    }
}

mod ziggurat_tables;

/// Sample a random number using the Ziggurat method (specifically the
/// ZIGNOR variant from Doornik 2005). Most of the arguments are
/// directly from the paper:
///
/// * `rng`: source of randomness
/// * `symmetric`: whether this is a symmetric distribution, or one-sided with P(x < 0) = 0.
/// * `X`: the $x_i$ abscissae.
/// * `F`: precomputed values of the PDF at the $x_i$, (i.e. $f(x_i)$)
/// * `F_DIFF`: precomputed values of $f(x_i) - f(x_{i+1})$
/// * `pdf`: the probability density function
/// * `zero_case`: manual sampling from the tail when we chose the
///    bottom box (i.e. i == 0)

// the perf improvement (25-50%) is definitely worth the extra code
// size from force-inlining.
#[inline(always)]
fn ziggurat<R:Rng>(
            rng: &mut R,
            symmetric: bool,
            x_tab: ziggurat_tables::ZigTable,
            f_tab: ziggurat_tables::ZigTable,
            pdf: |f64|: 'static -> f64,
            zero_case: |&mut R, f64|: 'static -> f64)
            -> f64 {
    static SCALE: f64 = (1u64 << 53) as f64;
    loop {
        // reimplement the f64 generation as an optimisation suggested
        // by the Doornik paper: we have a lot of precision-space
        // (i.e. there are 11 bits of the 64 of a u64 to use after
        // creating a f64), so we might as well reuse some to save
        // generating a whole extra random number. (Seems to be 15%
        // faster.)
        //
        // This unfortunately misses out on the benefits of direct
        // floating point generation if an RNG like dSMFT is
        // used. (That is, such RNGs create floats directly, highly
        // efficiently and overload next_f32/f64, so by not calling it
        // this may be slower than it would be otherwise.)
        // FIXME: investigate/optimise for the above.
        let bits: u64 = rng.gen();
        let i = (bits & 0xff) as uint;
        let f = (bits >> 11) as f64 / SCALE;

        // u is either U(-1, 1) or U(0, 1) depending on if this is a
        // symmetric distribution or not.
        let u = if symmetric {2.0 * f - 1.0} else {f};
        let x = u * x_tab[i];

        let test_x = if symmetric { x.abs() } else {x};

        // algebraically equivalent to |u| < x_tab[i+1]/x_tab[i] (or u < x_tab[i+1]/x_tab[i])
        if test_x < x_tab[i + 1] {
            return x;
        }
        if i == 0 {
            return zero_case(rng, u);
        }
        // algebraically equivalent to f1 + DRanU()*(f0 - f1) < 1
        if f_tab[i + 1] + (f_tab[i] - f_tab[i + 1]) * rng.gen() < pdf(x) {
            return x;
        }
    }
}

#[cfg(test)]
mod tests {
    use std::prelude::*;

    use {Rng, Rand};
    use super::{RandSample, WeightedChoice, Weighted, Sample, IndependentSample};

    #[deriving(PartialEq, Show)]
    struct ConstRand(uint);
    impl Rand for ConstRand {
        fn rand<R: Rng>(_: &mut R) -> ConstRand {
            ConstRand(0)
        }
    }

    // 0, 1, 2, 3, ...
    struct CountingRng { i: u32 }
    impl Rng for CountingRng {
        fn next_u32(&mut self) -> u32 {
            self.i += 1;
            self.i - 1
        }
        fn next_u64(&mut self) -> u64 {
            self.next_u32() as u64
        }
    }

    #[test]
    fn test_rand_sample() {
        let mut rand_sample = RandSample::<ConstRand>;

        assert_eq!(rand_sample.sample(&mut ::test::rng()), ConstRand(0));
        assert_eq!(rand_sample.ind_sample(&mut ::test::rng()), ConstRand(0));
    }
    #[test]
    fn test_weighted_choice() {
        // this makes assumptions about the internal implementation of
        // WeightedChoice, specifically: it doesn't reorder the items,
        // it doesn't do weird things to the RNG (so 0 maps to 0, 1 to
        // 1, internally; modulo a modulo operation).

        macro_rules! t (
            ($items:expr, $expected:expr) => {{
                let mut items = $items;
                let wc = WeightedChoice::new(items.as_mut_slice());
                let expected = $expected;

                let mut rng = CountingRng { i: 0 };

                for &val in expected.iter() {
                    assert_eq!(wc.ind_sample(&mut rng), val)
                }
            }}
        );

        t!(vec!(Weighted { weight: 1, item: 10i}), [10]);

        // skip some
        t!(vec!(Weighted { weight: 0, item: 20i},
                Weighted { weight: 2, item: 21i},
                Weighted { weight: 0, item: 22i},
                Weighted { weight: 1, item: 23i}),
           [21,21, 23]);

        // different weights
        t!(vec!(Weighted { weight: 4, item: 30i},
                Weighted { weight: 3, item: 31i}),
           [30,30,30,30, 31,31,31]);

        // check that we're binary searching
        // correctly with some vectors of odd
        // length.
        t!(vec!(Weighted { weight: 1, item: 40i},
                Weighted { weight: 1, item: 41i},
                Weighted { weight: 1, item: 42i},
                Weighted { weight: 1, item: 43i},
                Weighted { weight: 1, item: 44i}),
           [40, 41, 42, 43, 44]);
        t!(vec!(Weighted { weight: 1, item: 50i},
                Weighted { weight: 1, item: 51i},
                Weighted { weight: 1, item: 52i},
                Weighted { weight: 1, item: 53i},
                Weighted { weight: 1, item: 54i},
                Weighted { weight: 1, item: 55i},
                Weighted { weight: 1, item: 56i}),
           [50, 51, 52, 53, 54, 55, 56]);
    }

    #[test] #[should_fail]
    fn test_weighted_choice_no_items() {
        WeightedChoice::<int>::new([]);
    }
    #[test] #[should_fail]
    fn test_weighted_choice_zero_weight() {
        WeightedChoice::new(&mut [Weighted { weight: 0, item: 0i},
                                  Weighted { weight: 0, item: 1i}]);
    }
    #[test] #[should_fail]
    fn test_weighted_choice_weight_overflows() {
        let x = (-1) as uint / 2; // x + x + 2 is the overflow
        WeightedChoice::new(&mut [Weighted { weight: x, item: 0i },
                                  Weighted { weight: 1, item: 1i },
                                  Weighted { weight: x, item: 2i },
                                  Weighted { weight: 1, item: 3i }]);
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2014-2016 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Support for sending and receiving data link layer packets.

extern crate ipnetwork;
extern crate libc;
extern crate pnet_base;
extern crate pnet_sys;

#[cfg(feature = "serde")]
extern crate serde;
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};

use std::io;
use std::option::Option;
use std::time::Duration;

use ipnetwork::IpNetwork;

pub use pnet_base::{MacAddr, ParseMacAddrErr};

mod bindings;

#[cfg(windows)]
#[path = "winpcap.rs"]
mod backend;

#[cfg(windows)]
pub mod winpcap;

#[cfg(all(
    not(feature = "netmap"),
    any(target_os = "linux", target_os = "android")
))]
#[path = "linux.rs"]
mod backend;

#[cfg(any(target_os = "linux", target_os = "android"))]
pub mod linux;

#[cfg(all(
    not(feature = "netmap"),
    any(target_os = "freebsd", target_os = "openbsd", target_os = "netbsd", target_os = "macos", target_os = "ios")
))]
#[path = "bpf.rs"]
mod backend;

#[cfg(any(target_os = "freebsd", target_os = "netbsd", target_os = "macos", target_os = "ios"))]
pub mod bpf;

#[cfg(feature = "netmap")]
#[path = "netmap.rs"]
mod backend;

#[cfg(feature = "netmap")]
pub mod netmap;

#[cfg(feature = "pcap")]
pub mod pcap;

pub mod dummy;

/// Type alias for an `EtherType`.
pub type EtherType = u16;

/// Type of data link channel to present (Linux only).
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum ChannelType {
    /// Send and receive layer 2 packets directly, including headers.
    Layer2,
    /// Send and receive "cooked" packets - send and receive network layer packets.
    Layer3(EtherType),
}

/// A channel for sending and receiving at the data link layer.
///
/// NOTE: It is important to always include a catch-all variant in match statements using this
/// enum, since new variants may be added. For example:
///
/// ```ignore
/// match some_channel {
///     Ethernet(tx, rx) => { /* Handle Ethernet packets */ },
///     _ => panic!("Unhandled channel type")
/// }
/// ```
pub enum Channel {
    /// A datalink channel which sends and receives Ethernet packets.
    Ethernet(Box<dyn DataLinkSender>, Box<dyn DataLinkReceiver>),

    /// This variant should never be used.
    ///
    /// Including it allows new variants to be added to `Channel` without breaking existing code.
    PleaseIncludeACatchAllVariantWhenMatchingOnThisEnum,
}

/// Socket fanout type (Linux only).
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum FanoutType {
    HASH,
    LB,
    CPU,
    ROLLOVER,
    RND,
    QM,
    CBPF,
    EBPF,
}

/// Fanout settings (Linux only).
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct FanoutOption {
    pub group_id: u16,
    pub fanout_type: FanoutType,
    pub defrag: bool,
    pub rollover: bool,
}

/// A generic configuration type, encapsulating all options supported by each backend.
///
/// Each option should be treated as a hint - each backend is free to ignore any and all
/// options which don't apply to it.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct Config {
    /// The size of buffer to use when writing packets. Defaults to 4096.
    pub write_buffer_size: usize,

    /// The size of buffer to use when reading packets. Defaults to 4096.
    pub read_buffer_size: usize,

    /// Linux/BPF/Netmap only: The read timeout. Defaults to None.
    pub read_timeout: Option<Duration>,

    /// Linux/BPF/Netmap only: The write timeout. Defaults to None.
    pub write_timeout: Option<Duration>,

    /// Linux only: Specifies whether to read packets at the datalink layer or network layer.
    /// Defaults to Layer2
    pub channel_type: ChannelType,

    /// BPF/OS X only: The number of /dev/bpf* file descriptors to attempt before failing. Defaults
    /// to: 1000.
    pub bpf_fd_attempts: usize,

    pub linux_fanout: Option<FanoutOption>,

    pub promiscuous: bool,
}

impl Default for Config {
    fn default() -> Config {
        Config {
            write_buffer_size: 4096,
            read_buffer_size: 4096,
            read_timeout: None,
            write_timeout: None,
            channel_type: ChannelType::Layer2,
            bpf_fd_attempts: 1000,
            linux_fanout: None,
            promiscuous: true,
        }
    }
}

/// Create a new datalink channel for sending and receiving data.
///
/// This allows for sending and receiving packets at the data link layer.
///
/// A list of network interfaces can be retrieved using datalink::interfaces().
///
/// The configuration serves as a hint to the backend - some or all of it may be used or ignored,
/// depending on which backend is used.
///
/// When matching on the returned channel, make sure to include a catch-all so that code doesn't
/// break when new channel types are added.
#[inline]
pub fn channel(network_interface: &NetworkInterface, configuration: Config) -> io::Result<Channel> {
    backend::channel(network_interface, (&configuration).into())
}

/// Trait to enable sending `$packet` packets.
pub trait DataLinkSender: Send {
    /// Create and send a number of packets.
    ///
    /// This will call `func` `num_packets` times. The function will be provided with a
    /// mutable packet to manipulate, which will then be sent. This allows packets to be
    /// built in-place, avoiding the copy required for `send`. If there is not sufficient
    /// capacity in the buffer, None will be returned.
    fn build_and_send(
        &mut self,
        num_packets: usize,
        packet_size: usize,
        func: &mut dyn FnMut(&mut [u8]),
    ) -> Option<io::Result<()>>;

    /// Send a packet.
    ///
    /// This may require an additional copy compared to `build_and_send`, depending on the
    /// operating system being used. The second parameter is currently ignored, however
    /// `None` should be passed.
    fn send_to(&mut self, packet: &[u8], dst: Option<NetworkInterface>) -> Option<io::Result<()>>;
}

/// Structure for receiving packets at the data link layer. Should be constructed using
/// `datalink_channel()`.
pub trait DataLinkReceiver: Send {
    /// Get the next ethernet frame in the channel.
    fn next(&mut self) -> io::Result<&[u8]>;
}

/// Represents a network interface and its associated addresses.
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[derive(Clone, PartialEq, Eq, Debug, Hash)]
pub struct NetworkInterface {
    /// The name of the interface.
    pub name: String,
    /// A description of the interface.
    pub description: String,
    /// The interface index (operating system specific).
    pub index: u32,
    /// A MAC address for the interface.
    pub mac: Option<MacAddr>,
    /// IP addresses and netmasks for the interface.
    pub ips: Vec<IpNetwork>,
    /// Operating system specific flags for the interface.
    pub flags: u32,
}

impl NetworkInterface {
    /// Retrieve the MAC address associated with the interface.
    #[deprecated(
        since = "0.26.0",
        note = "Please use NetworkInterface's field 'mac' instead."
    )]
    pub fn mac_address(&self) -> MacAddr {
        self.mac.unwrap()
    }

    pub fn is_up(&self) -> bool {
        self.flags & (pnet_sys::IFF_UP as u32) != 0
    }
    pub fn is_broadcast(&self) -> bool {
        self.flags & (pnet_sys::IFF_BROADCAST as u32) != 0
    }
    /// Is the interface a loopback interface?
    pub fn is_loopback(&self) -> bool {
        self.flags & (pnet_sys::IFF_LOOPBACK as u32) != 0
    }
    pub fn is_point_to_point(&self) -> bool {
        self.flags & (pnet_sys::IFF_POINTOPOINT as u32) != 0
    }
    pub fn is_multicast(&self) -> bool {
        self.flags & (pnet_sys::IFF_MULTICAST as u32) != 0
    }
}

impl ::std::fmt::Display for NetworkInterface {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        const FLAGS: [&'static str; 5] =
            ["UP", "BROADCAST", "LOOPBACK", "POINTOPOINT", "MULTICAST"];
        let flags = if self.flags > 0 {
            let rets = [
                self.is_up(),
                self.is_broadcast(),
                self.is_loopback(),
                self.is_point_to_point(),
                self.is_multicast(),
            ];
            format!(
                "{:X}<{}>",
                self.flags,
                rets.iter()
                    .zip(FLAGS.iter())
                    .filter(|&(ret, _)| ret == &true)
                    .map(|(_, name)| name.to_string())
                    .collect::<Vec<String>>()
                    .join(",")
            )
        } else {
            format!("{:X}", self.flags)
        };

        let mac = self
            .mac
            .map(|mac| mac.to_string())
            .unwrap_or("N/A".to_owned());
        let ips = if self.ips.len() > 0 {
            format!(
                "\n{}",
                self.ips
                    .iter()
                    .map(|ip| {
                        if ip.is_ipv4() {
                            format!("       inet: {}", ip)
                        } else {
                            format!("      inet6: {}", ip)
                        }
                    })
                    .collect::<Vec<String>>()
                    .join("\n")
            )
        } else {
            "".to_string()
        };

        write!(
            f,
            "{}: flags={}
      index: {}
      ether: {}{}",
            self.name, flags, self.index, mac, ips
        )
    }
}

/// Get a list of available network interfaces for the current machine.
///
/// If you need the default network interface, you can choose the first
/// one that is up, not loopback and has an IP. This is not guaranteed to
/// work on each system but should work for basic packet sniffing:
///
/// ```
/// use pnet_datalink::interfaces;
///
/// // Get a vector with all network interfaces found
/// let all_interfaces = interfaces();
///
/// // Search for the default interface - the one that is
/// // up, not loopback and has an IP.
/// let default_interface = all_interfaces
///     .iter()
///     .find(|e| e.is_up() && !e.is_loopback() && !e.ips.is_empty());
///
/// match default_interface {
///     Some(interface) => println!("Found default interface with [{}].", interface.name),
///     None => println!("Error while finding the default interface."),
/// }
///
/// ```
///
pub fn interfaces() -> Vec<NetworkInterface> {
    backend::interfaces()
}
<commit_msg>module path should be pnet::datalink<commit_after>// Copyright (c) 2014-2016 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Support for sending and receiving data link layer packets.

extern crate ipnetwork;
extern crate libc;
extern crate pnet_base;
extern crate pnet_sys;

#[cfg(feature = "serde")]
extern crate serde;
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};

use std::io;
use std::option::Option;
use std::time::Duration;

use ipnetwork::IpNetwork;

pub use pnet_base::{MacAddr, ParseMacAddrErr};

mod bindings;

#[cfg(windows)]
#[path = "winpcap.rs"]
mod backend;

#[cfg(windows)]
pub mod winpcap;

#[cfg(all(
    not(feature = "netmap"),
    any(target_os = "linux", target_os = "android")
))]
#[path = "linux.rs"]
mod backend;

#[cfg(any(target_os = "linux", target_os = "android"))]
pub mod linux;

#[cfg(all(
    not(feature = "netmap"),
    any(target_os = "freebsd", target_os = "openbsd", target_os = "netbsd", target_os = "macos", target_os = "ios")
))]
#[path = "bpf.rs"]
mod backend;

#[cfg(any(target_os = "freebsd", target_os = "netbsd", target_os = "macos", target_os = "ios"))]
pub mod bpf;

#[cfg(feature = "netmap")]
#[path = "netmap.rs"]
mod backend;

#[cfg(feature = "netmap")]
pub mod netmap;

#[cfg(feature = "pcap")]
pub mod pcap;

pub mod dummy;

/// Type alias for an `EtherType`.
pub type EtherType = u16;

/// Type of data link channel to present (Linux only).
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum ChannelType {
    /// Send and receive layer 2 packets directly, including headers.
    Layer2,
    /// Send and receive "cooked" packets - send and receive network layer packets.
    Layer3(EtherType),
}

/// A channel for sending and receiving at the data link layer.
///
/// NOTE: It is important to always include a catch-all variant in match statements using this
/// enum, since new variants may be added. For example:
///
/// ```ignore
/// match some_channel {
///     Ethernet(tx, rx) => { /* Handle Ethernet packets */ },
///     _ => panic!("Unhandled channel type")
/// }
/// ```
pub enum Channel {
    /// A datalink channel which sends and receives Ethernet packets.
    Ethernet(Box<dyn DataLinkSender>, Box<dyn DataLinkReceiver>),

    /// This variant should never be used.
    ///
    /// Including it allows new variants to be added to `Channel` without breaking existing code.
    PleaseIncludeACatchAllVariantWhenMatchingOnThisEnum,
}

/// Socket fanout type (Linux only).
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum FanoutType {
    HASH,
    LB,
    CPU,
    ROLLOVER,
    RND,
    QM,
    CBPF,
    EBPF,
}

/// Fanout settings (Linux only).
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct FanoutOption {
    pub group_id: u16,
    pub fanout_type: FanoutType,
    pub defrag: bool,
    pub rollover: bool,
}

/// A generic configuration type, encapsulating all options supported by each backend.
///
/// Each option should be treated as a hint - each backend is free to ignore any and all
/// options which don't apply to it.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct Config {
    /// The size of buffer to use when writing packets. Defaults to 4096.
    pub write_buffer_size: usize,

    /// The size of buffer to use when reading packets. Defaults to 4096.
    pub read_buffer_size: usize,

    /// Linux/BPF/Netmap only: The read timeout. Defaults to None.
    pub read_timeout: Option<Duration>,

    /// Linux/BPF/Netmap only: The write timeout. Defaults to None.
    pub write_timeout: Option<Duration>,

    /// Linux only: Specifies whether to read packets at the datalink layer or network layer.
    /// Defaults to Layer2
    pub channel_type: ChannelType,

    /// BPF/OS X only: The number of /dev/bpf* file descriptors to attempt before failing. Defaults
    /// to: 1000.
    pub bpf_fd_attempts: usize,

    pub linux_fanout: Option<FanoutOption>,

    pub promiscuous: bool,
}

impl Default for Config {
    fn default() -> Config {
        Config {
            write_buffer_size: 4096,
            read_buffer_size: 4096,
            read_timeout: None,
            write_timeout: None,
            channel_type: ChannelType::Layer2,
            bpf_fd_attempts: 1000,
            linux_fanout: None,
            promiscuous: true,
        }
    }
}

/// Create a new datalink channel for sending and receiving data.
///
/// This allows for sending and receiving packets at the data link layer.
///
/// A list of network interfaces can be retrieved using datalink::interfaces().
///
/// The configuration serves as a hint to the backend - some or all of it may be used or ignored,
/// depending on which backend is used.
///
/// When matching on the returned channel, make sure to include a catch-all so that code doesn't
/// break when new channel types are added.
#[inline]
pub fn channel(network_interface: &NetworkInterface, configuration: Config) -> io::Result<Channel> {
    backend::channel(network_interface, (&configuration).into())
}

/// Trait to enable sending `$packet` packets.
pub trait DataLinkSender: Send {
    /// Create and send a number of packets.
    ///
    /// This will call `func` `num_packets` times. The function will be provided with a
    /// mutable packet to manipulate, which will then be sent. This allows packets to be
    /// built in-place, avoiding the copy required for `send`. If there is not sufficient
    /// capacity in the buffer, None will be returned.
    fn build_and_send(
        &mut self,
        num_packets: usize,
        packet_size: usize,
        func: &mut dyn FnMut(&mut [u8]),
    ) -> Option<io::Result<()>>;

    /// Send a packet.
    ///
    /// This may require an additional copy compared to `build_and_send`, depending on the
    /// operating system being used. The second parameter is currently ignored, however
    /// `None` should be passed.
    fn send_to(&mut self, packet: &[u8], dst: Option<NetworkInterface>) -> Option<io::Result<()>>;
}

/// Structure for receiving packets at the data link layer. Should be constructed using
/// `datalink_channel()`.
pub trait DataLinkReceiver: Send {
    /// Get the next ethernet frame in the channel.
    fn next(&mut self) -> io::Result<&[u8]>;
}

/// Represents a network interface and its associated addresses.
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[derive(Clone, PartialEq, Eq, Debug, Hash)]
pub struct NetworkInterface {
    /// The name of the interface.
    pub name: String,
    /// A description of the interface.
    pub description: String,
    /// The interface index (operating system specific).
    pub index: u32,
    /// A MAC address for the interface.
    pub mac: Option<MacAddr>,
    /// IP addresses and netmasks for the interface.
    pub ips: Vec<IpNetwork>,
    /// Operating system specific flags for the interface.
    pub flags: u32,
}

impl NetworkInterface {
    /// Retrieve the MAC address associated with the interface.
    #[deprecated(
        since = "0.26.0",
        note = "Please use NetworkInterface's field 'mac' instead."
    )]
    pub fn mac_address(&self) -> MacAddr {
        self.mac.unwrap()
    }

    pub fn is_up(&self) -> bool {
        self.flags & (pnet_sys::IFF_UP as u32) != 0
    }
    pub fn is_broadcast(&self) -> bool {
        self.flags & (pnet_sys::IFF_BROADCAST as u32) != 0
    }
    /// Is the interface a loopback interface?
    pub fn is_loopback(&self) -> bool {
        self.flags & (pnet_sys::IFF_LOOPBACK as u32) != 0
    }
    pub fn is_point_to_point(&self) -> bool {
        self.flags & (pnet_sys::IFF_POINTOPOINT as u32) != 0
    }
    pub fn is_multicast(&self) -> bool {
        self.flags & (pnet_sys::IFF_MULTICAST as u32) != 0
    }
}

impl ::std::fmt::Display for NetworkInterface {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        const FLAGS: [&'static str; 5] =
            ["UP", "BROADCAST", "LOOPBACK", "POINTOPOINT", "MULTICAST"];
        let flags = if self.flags > 0 {
            let rets = [
                self.is_up(),
                self.is_broadcast(),
                self.is_loopback(),
                self.is_point_to_point(),
                self.is_multicast(),
            ];
            format!(
                "{:X}<{}>",
                self.flags,
                rets.iter()
                    .zip(FLAGS.iter())
                    .filter(|&(ret, _)| ret == &true)
                    .map(|(_, name)| name.to_string())
                    .collect::<Vec<String>>()
                    .join(",")
            )
        } else {
            format!("{:X}", self.flags)
        };

        let mac = self
            .mac
            .map(|mac| mac.to_string())
            .unwrap_or("N/A".to_owned());
        let ips = if self.ips.len() > 0 {
            format!(
                "\n{}",
                self.ips
                    .iter()
                    .map(|ip| {
                        if ip.is_ipv4() {
                            format!("       inet: {}", ip)
                        } else {
                            format!("      inet6: {}", ip)
                        }
                    })
                    .collect::<Vec<String>>()
                    .join("\n")
            )
        } else {
            "".to_string()
        };

        write!(
            f,
            "{}: flags={}
      index: {}
      ether: {}{}",
            self.name, flags, self.index, mac, ips
        )
    }
}

/// Get a list of available network interfaces for the current machine.
///
/// If you need the default network interface, you can choose the first
/// one that is up, not loopback and has an IP. This is not guaranteed to
/// work on each system but should work for basic packet sniffing:
///
/// ```
/// use pnet::datalink::interfaces;
///
/// // Get a vector with all network interfaces found
/// let all_interfaces = interfaces();
///
/// // Search for the default interface - the one that is
/// // up, not loopback and has an IP.
/// let default_interface = all_interfaces
///     .iter()
///     .find(|e| e.is_up() && !e.is_loopback() && !e.ips.is_empty());
///
/// match default_interface {
///     Some(interface) => println!("Found default interface with [{}].", interface.name),
///     None => println!("Error while finding the default interface."),
/// }
///
/// ```
///
pub fn interfaces() -> Vec<NetworkInterface> {
    backend::interfaces()
}
<|endoftext|>