<commit_before>//! Check running processes

use std::collections::HashSet;
use std::str::FromStr;

use log::LevelFilter::{Debug, Trace, Warn};
use log::{debug, trace};
use nix::sys::signal::{kill, Signal as NixSignal};
use nix::unistd::{getpid, getppid, Pid};
use regex::Regex;
use structopt::StructOpt;

use tabin_plugins::procfs::pid::{Process, State};
use tabin_plugins::procfs::{LoadProcsError, ProcFsError, ProcMap, RunningProcs};
use tabin_plugins::Status;

const LOG_VAR: &str = "TABIN_LOG";

/// Check that an expected number of processes are running.
///
/// Optionally, kill unwanted processes.
#[derive(StructOpt, Debug)]
#[structopt(
    name = "check-procs (part of tabin-plugins)",
    raw(setting = "structopt::clap::AppSettings::ColoredHelp"),
    after_help = "Examples:

    Ensure at least two nginx processes are running:

        check-procs --crit-under 2 nginx

    Ensure there are not more than 30 zombie proccesses on the system:

        check-procs --crit-over 30 --state zombie

    Ensure that there are not more than 5 java processes running MyMainClass
    that are in the zombie *or* waiting states. Note that since there can be
    multiple states the regex must come before the `state` flag:

        check-procs 'java.*MyMainClass' --crit-over 5 --state zombie waiting

    Ensure that there are at least three (running or waiting) (cassandra or
    postgres) processes:

        check-procs --crit-under 3 --state=running --state=waiting 'cassandra|postgres'"
)]
struct Args {
    #[structopt(help = "Regex that command and its arguments must match")]
    pattern: Option<Regex>,
    #[structopt(
        long = "crit-under",
        name = "N",
        help = "Error if there are fewer than <N> procs matching <pattern>"
    )]
    crit_under: Option<usize>,
    #[structopt(
        long = "crit-over",
        name = "M",
        help = "Error if there are more than <M> procs matching <pattern>"
    )]
    crit_over: Option<usize>,

    #[structopt(
        long = "state",
        help = "Filter to only processes in these states. \
                If passed multiple times, processes matching any state are included.\n\
                Choices: running sleeping uninterruptible-sleep waiting stopped zombie"
    )]
    states: Vec<State>,

    #[structopt(
        long = "allow-unparseable-procs",
        help = "In combination with --crit-over M this will not alert if any \
                processes cannot be parsed"
    )]
    allow_unparseable_procs: bool,

    #[structopt(
        long = "kill-matching",
        name = "SIGNAL",
        help = "If *any* processes match, then kill them with the provided signal \
                which can be either an integer or a name like KILL or SIGTERM. \
                This option does not affect the exit status, all matches are always \
                killed, and if --crit-under/over are violated then then this will \
                still exit critical."
    )]
    kill_matching: Option<Signal>,

    #[structopt(
        long = "kill-parents-of-matching",
        name = "PARENT_SIGNAL",
        help = "If *any* processes match, then kill their parents with the provided \
                signal which can be either an integer or a name like KILL or SIGTERM. \
                This has the same exit status behavior as kill-matching."
    )]
    kill_matching_parents: Option<Signal>,

    /// print debug logs, use multiple times to make it more verbose
    #[structopt(short = "v", long = "verbose", parse(from_occurrences))]
    verbose: u8,
}

/// Our own signal wrapper so that we can implement a forgiving FromStr for `nix::sys::Signal`
#[derive(Debug)]
struct Signal(NixSignal);

impl FromStr for Signal {
    type Err = String;
    fn from_str(s: &str) -> Result<Signal, String> {
        let sig: Result<i32, _> = s.parse();
        match sig {
            Ok(integer) => {
                Ok(Signal(NixSignal::from_c_int(integer).map_err(|_| {
                    format!("Not a valid signal integer: {}", s)
                })?))
            }
            Err(_) => Ok(Signal(match s {
                "SIGHUP" | "HUP" => NixSignal::SIGHUP,
                "SIGINT" | "INT" => NixSignal::SIGINT,
                "SIGQUIT" | "QUIT" => NixSignal::SIGQUIT,
                "SIGILL" | "ILL" => NixSignal::SIGILL,
                "SIGTRAP" | "TRAP" => NixSignal::SIGTRAP,
                "SIGABRT" | "ABRT" => NixSignal::SIGABRT,
                "SIGBUS" | "BUS" => NixSignal::SIGBUS,
                "SIGFPE" | "FPE" => NixSignal::SIGFPE,
                "SIGKILL" | "KILL" => NixSignal::SIGKILL,
                "SIGUSR1" | "USR1" => NixSignal::SIGUSR1,
                "SIGSEGV" | "SEGV" => NixSignal::SIGSEGV,
                "SIGUSR2" | "USR2" => NixSignal::SIGUSR2,
                "SIGPIPE" | "PIPE" => NixSignal::SIGPIPE,
                "SIGALRM" | "ALRM" => NixSignal::SIGALRM,
                "SIGTERM" | "TERM" => NixSignal::SIGTERM,
                _ => return Err(format!("Could not parse {:?} as an int or named signal", s)),
            })),
        }
    }
}

fn parse_args() -> Args {
    let args = Args::from_args();
    if args.pattern.is_none() && args.states.is_empty() {
        println!("At least one of a pattern or some states are required for this to do anything");
        Status::Critical.exit();
    }
    if let (None, None) = (args.crit_under, args.crit_over) {
        println!("At least one of --crit-under or --crit-over must be provided");
        Status::Critical.exit();
    }
    args
}

fn main() {
    let args = parse_args();
    env_logger::Builder::from_env(LOG_VAR)
        .filter_level(match args.verbose {
            0 => Warn,
            1 => Debug,
            _ => Trace,
        })
        .init();

    let should_die = if let Some(_) = args.crit_over {
        !args.allow_unparseable_procs
    } else {
        false
    };
    let procs = load_procs(should_die);
    let re = args.pattern.as_ref().map(|r| r.to_string());

    let matches = filter_procs(args.pattern, &args.states, &procs.0);

    let mut status = Status::Ok;
    if let Some(crit_over) = args.crit_over {
        if matches.len() > crit_over {
            status = Status::Critical;
            println!(
                "CRITICAL: there are {} process that match {:?} (greater than {})",
                matches.len(),
                re,
                crit_over
            );
        }
    };
    if let Some(crit_under) = args.crit_under {
        if matches.len() < crit_under {
            status = Status::Critical;
            println!(
                "CRITICAL: there are {} process that match {:?} (less than {})",
                matches.len(),
                re,
                crit_under
            );
        }
    }

    if status == Status::Ok {
        match (args.crit_over, args.crit_under) {
            (Some(o), Some(u)) => println!(
                "OKAY: There are {} matching procs (between {} and {})",
                matches.len(),
                o,
                u
            ),
            (Some(o), None) => println!(
                "OKAY: There are {} matching procs (less than {})",
                matches.len(),
                o
            ),
            (None, Some(u)) => println!(
                "OKAY: There are {} matching procs (greater than {})",
                matches.len(),
                u
            ),
            (None, None) => unreachable!(),
        }
    }
    if matches.len() > 0 {
        println!("INFO: Matching processes:");
        for process in matches.iter().take(20) {
            println!("[{:>5}] {}", process.0, process.1.useful_cmdline());
        }
        if matches.len() > 20 {
            println!("And {} more...", matches.len() - 20)
        }
    }

    if args.kill_matching.is_some() {
        let signal = args.kill_matching.unwrap().0;
        let errors: Vec<_> = matches
            .iter()
            .map(|&(pid, _)| kill(*pid, signal))
            .filter(|result| result.is_err())
            .collect();
        if !errors.is_empty() {
            println!("INFO: There were {} errors killing processes", errors.len());
        }
    }
    if args.kill_matching_parents.is_some() {
        let signal = args.kill_matching_parents.unwrap().0;
        let mut parents = HashSet::new();
        let errors: Vec<_> = matches
            .iter()
            .map(|&(_, process)| process.stat.ppid)
            .filter(|ppid| parents.insert(*ppid))
            .map(|ppid| kill(ppid, signal))
            .filter(|result| result.is_err())
            .collect();
        if !errors.is_empty() {
            println!("INFO: There were {} errors killing processes", errors.len());
        }
    }

    status.exit();
}

/// Load currently running procs, and die if there is a surprising error
///
/// Normally if this can load *any* processes it returns what it can find, and
/// prints errors for procs that can't be parsed. But if `die_on_any_errors` is
/// true it dies if it cannot parse a *single* process.
fn load_procs(die_on_any_errors: bool) -> RunningProcs {
    match RunningProcs::currently_running() {
        Ok(procs) => procs,
        Err(ProcFsError::LoadProcsError(LoadProcsError { procs, errors })) => {
            let mut saw_real_error = false;
            for err in errors {
                match err {
                    // If the process no longer exists, that's fine
                    ProcFsError::Io(_) => {}
                    err => {
                        saw_real_error = true;
                        println!("WARN: Unexpected error loading some processes: {}", err);
                    }
                }
            }
            if die_on_any_errors && saw_real_error {
                Status::Critical.exit();
            }
            procs
        }
        Err(err) => {
            println!("ERROR: unable to load processes: {}", err);
            Status::Critical.exit();
        }
    }
}

/// Filter to processes that match the condition
///
/// The condition is:
///
/// * Does not match me or my parent process
/// * Does match the regex, if it is present
/// * Does match *any* of the states
fn filter_procs<'m>(
    re: Option<Regex>,
    states: &[State],
    procs: &'m ProcMap,
) -> Vec<(&'m Pid, &'m Process)> {
    let me = getpid();
    let parent = getppid();

    debug!("filtering for procs that match: {:?} {:?}", re, states);

    procs
        .iter()
        .filter(|&(pid, _process)| *pid != me && *pid != parent)
        .filter(|&(_pid, process)| states.is_empty() || states.contains(&process.stat.state))
        .filter(|&(_pid, process)| {
            re.as_ref()
                .map(|re| {
                    let cmdline = process.useful_cmdline();
                    let is_match = re.is_match(&cmdline);
                    trace!("is_match={} process={}", is_match, cmdline);
                    is_match
                })
                .unwrap_or(true)
        })
        .collect()
}

#[cfg(test)]
mod unit {
    use structopt::StructOpt;

    use tabin_plugins::procfs::pid::Process;

    use super::*;

    #[test]
    fn validate_argparse() {
        let args = Args::from_iter(["c-p", "some.*proc", "--crit-under=1"].into_iter());
        assert_eq!(args.crit_under, Some(1));
    }

    #[test]
    fn validate_parse_zombies() {
        let args = Args::from_iter(["c-p", "some.*proc", "--state=zombie"].into_iter());
        assert_eq!(args.states, [State::Zombie]);
        let args =
            Args::from_iter(["c-p", "some.*proc", "--state=zombie", "--state", "S"].into_iter());
        assert_eq!(args.states, [State::Zombie, State::Sleeping]);
    }

    #[test]
    fn validate_parse_zombies_and_pattern() {
        let args = Args::from_iter(["c-p", "--state", "zombie", "--", "some.*proc"].into_iter());
        assert_eq!(args.states, [State::Zombie]);
        let args = Args::from_iter(["c-p", "--state=zombie", "some.*proc"].into_iter());
        assert_eq!(args.states, [State::Zombie]);
        let args = Args::from_iter(["c-p", "so.*proc", "--state", "zombie", "waiting"].into_iter());
        assert_eq!(args.states, [State::Zombie, State::Waiting]);
    }

    // Waiting for structopt 0.2.8 to be released with the from_iter_safe method
    // #[test]
    // #[should_panic]
    // fn regexes_must_be_valid() {
    //     Args::from_iter_safe(["c-p", "--crit-over=5", "["].into_iter()).unwrap();
    // }

    #[test]
    fn filter_procs_handles_patterns() {
        let mut procs = vec![Process::default(); 5];
        procs[2]
            .cmdline
            .raw
            .append(&mut vec!["hello".into(), "jar".into()]);
        let proc_map = vec_to_procmap(procs);
        let filtered = filter_procs(regex("llo.*ar"), &[], &proc_map);
        assert_eq!(filtered.len(), 1);
    }

    #[test]
    fn filter_procs_handles_single_state() {
        let mut procs = vec![Process::default(); 5];
        procs[2].stat.state = State::Zombie;
        let proc_map = vec_to_procmap(procs);
        let filtered = filter_procs(None, &[State::Zombie], &proc_map);
        assert_eq!(filtered.len(), 1);
    }

    #[test]
    fn filter_procs_handles_multiple_states() {
        let mut procs = vec![Process::default(); 5];
        procs[2].stat.state = State::Zombie;
        procs[3].stat.state = State::Stopped;
        let proc_map = vec_to_procmap(procs);

        let filtered = filter_procs(None, &[State::Zombie], &proc_map);
        assert_eq!(filtered.len(), 1);

        let filtered = filter_procs(None, &[State::Zombie, State::Stopped], &proc_map);
        assert_eq!(filtered.len(), 2);
    }

    #[test]
    fn filter_procs_handles_multiple_states_and_regex() {
        let mut procs = vec![Process::default(); 5];
        procs[2].stat.state = State::Zombie;
        procs[2].cmdline.raw.push("example".into());

        procs[3].stat.state = State::Stopped;
        let proc_map = vec_to_procmap(procs.clone());

        let filtered = filter_procs(regex("exa"), &[], &proc_map);
        assert_eq!(filtered.len(), 1);

        let filtered = filter_procs(regex("exa"), &[State::Zombie, State::Stopped], &proc_map);
        assert_eq!(filtered.len(), 1);
        assert_eq!(filtered[0].1, &procs[2]);

        let filtered = filter_procs(regex("exa"), &[State::Stopped], &proc_map);
        assert_eq!(filtered.len(), 0);
    }

    fn regex(re: &str) -> Option<Regex> {
        Some(Regex::new(re).unwrap())
    }

    fn vec_to_procmap(procs: Vec<Process>) -> ProcMap {
        procs
            .into_iter()
            .enumerate()
            .map(|(i, process)| (Pid::from_raw(i as i32), process))
            .collect()
    }
}
<commit_msg>[check-procs] improve output<commit_after>//! Check running processes

use std::collections::HashSet;
use std::str::FromStr;

use log::LevelFilter::{Debug, Trace, Warn};
use log::{debug, trace};
use nix::sys::signal::{kill, Signal as NixSignal};
use nix::unistd::{getpid, getppid, Pid};
use regex::Regex;
use structopt::StructOpt;

use tabin_plugins::procfs::pid::{Process, State};
use tabin_plugins::procfs::{LoadProcsError, ProcFsError, ProcMap, RunningProcs};
use tabin_plugins::Status;

const LOG_VAR: &str = "TABIN_LOG";

/// Check that an expected number of processes are running.
///
/// Optionally, kill unwanted processes.
#[derive(StructOpt, Debug)]
#[structopt(
    name = "check-procs (part of tabin-plugins)",
    raw(setting = "structopt::clap::AppSettings::ColoredHelp"),
    after_help = "Examples:

    Ensure at least two nginx processes are running:

        check-procs --crit-under 2 nginx

    Ensure there are not more than 30 zombie proccesses on the system:

        check-procs --crit-over 30 --state zombie

    Ensure that there are not more than 5 java processes running MyMainClass
    that are in the zombie *or* waiting states. Note that since there can be
    multiple states the regex must come before the `state` flag:

        check-procs 'java.*MyMainClass' --crit-over 5 --state zombie waiting

    Ensure that there are at least three (running or waiting) (cassandra or
    postgres) processes:

        check-procs --crit-under 3 --state=running --state=waiting 'cassandra|postgres'"
)]
struct Args {
    #[structopt(help = "Regex that command and its arguments must match")]
    pattern: Option<Regex>,
    #[structopt(
        long = "crit-under",
        name = "N",
        help = "Error if there are fewer than <N> procs matching <pattern>"
    )]
    crit_under: Option<usize>,
    #[structopt(
        long = "crit-over",
        name = "M",
        help = "Error if there are more than <M> procs matching <pattern>"
    )]
    crit_over: Option<usize>,

    #[structopt(
        long = "state",
        help = "Filter to only processes in these states. \
                If passed multiple times, processes matching any state are included.\n\
                Choices: running sleeping uninterruptible-sleep waiting stopped zombie"
    )]
    states: Vec<State>,

    #[structopt(
        long = "allow-unparseable-procs",
        help = "In combination with --crit-over M this will not alert if any \
                processes cannot be parsed"
    )]
    allow_unparseable_procs: bool,

    #[structopt(
        long = "kill-matching",
        name = "SIGNAL",
        help = "If *any* processes match, then kill them with the provided signal \
                which can be either an integer or a name like KILL or SIGTERM. \
                This option does not affect the exit status, all matches are always \
                killed, and if --crit-under/over are violated then then this will \
                still exit critical."
    )]
    kill_matching: Option<Signal>,

    #[structopt(
        long = "kill-parents-of-matching",
        name = "PARENT_SIGNAL",
        help = "If *any* processes match, then kill their parents with the provided \
                signal which can be either an integer or a name like KILL or SIGTERM. \
                This has the same exit status behavior as kill-matching."
    )]
    kill_matching_parents: Option<Signal>,

    /// print debug logs, use multiple times to make it more verbose
    #[structopt(short = "v", long = "verbose", parse(from_occurrences))]
    verbose: u8,
}

/// Our own signal wrapper so that we can implement a forgiving FromStr for `nix::sys::Signal`
#[derive(Debug)]
struct Signal(NixSignal);

impl FromStr for Signal {
    type Err = String;
    fn from_str(s: &str) -> Result<Signal, String> {
        let sig: Result<i32, _> = s.parse();
        match sig {
            Ok(integer) => {
                Ok(Signal(NixSignal::from_c_int(integer).map_err(|_| {
                    format!("Not a valid signal integer: {}", s)
                })?))
            }
            Err(_) => Ok(Signal(match s {
                "SIGHUP" | "HUP" => NixSignal::SIGHUP,
                "SIGINT" | "INT" => NixSignal::SIGINT,
                "SIGQUIT" | "QUIT" => NixSignal::SIGQUIT,
                "SIGILL" | "ILL" => NixSignal::SIGILL,
                "SIGTRAP" | "TRAP" => NixSignal::SIGTRAP,
                "SIGABRT" | "ABRT" => NixSignal::SIGABRT,
                "SIGBUS" | "BUS" => NixSignal::SIGBUS,
                "SIGFPE" | "FPE" => NixSignal::SIGFPE,
                "SIGKILL" | "KILL" => NixSignal::SIGKILL,
                "SIGUSR1" | "USR1" => NixSignal::SIGUSR1,
                "SIGSEGV" | "SEGV" => NixSignal::SIGSEGV,
                "SIGUSR2" | "USR2" => NixSignal::SIGUSR2,
                "SIGPIPE" | "PIPE" => NixSignal::SIGPIPE,
                "SIGALRM" | "ALRM" => NixSignal::SIGALRM,
                "SIGTERM" | "TERM" => NixSignal::SIGTERM,
                _ => return Err(format!("Could not parse {:?} as an int or named signal", s)),
            })),
        }
    }
}

fn parse_args() -> Args {
    let args = Args::from_args();
    if args.pattern.is_none() && args.states.is_empty() {
        println!("At least one of a pattern or some states are required for this to do anything");
        Status::Critical.exit();
    }
    if let (None, None) = (args.crit_under, args.crit_over) {
        println!("At least one of --crit-under or --crit-over must be provided");
        Status::Critical.exit();
    }
    args
}

fn main() {
    let args = parse_args();
    env_logger::Builder::from_env(LOG_VAR)
        .filter_level(match args.verbose {
            0 => Warn,
            1 => Debug,
            _ => Trace,
        })
        .init();

    let should_die = if let Some(_) = args.crit_over {
        !args.allow_unparseable_procs
    } else {
        false
    };
    let procs = load_procs(should_die);
    let re_ = args.pattern.as_ref().map(|s| s.to_string());
    let re = || re_.as_ref().map(|s| &**s).unwrap_or("<ANYTHING>");

    let matches = filter_procs(&args.pattern, &args.states, &procs.0);

    let mut status = Status::Ok;
    if let Some(crit_over) = args.crit_over {
        if matches.len() > crit_over {
            status = Status::Critical;
        }
    };
    if let Some(crit_under) = args.crit_under {
        if matches.len() < crit_under {
            status = Status::Critical;
        }
    }

    print!("{}: there are {} procs ", status, matches.len());
    if args.pattern.is_some() {
        print!("that match '{}' ", re());
    }
    if !args.states.is_empty() {
        print!("with any state in {:?} ", args.states);
    }
    let prefix = if status == Status::Critical {
        "not "
    } else {
        ""
    };
    match (args.crit_over, args.crit_under) {
        (Some(over), Some(under)) => println!("({}between {} and {})", prefix, over, under),
        (Some(crit_over), None) => println!("({}less or equal to {})", prefix, crit_over),
        (None, Some(crit_under)) => println!("({}greater or equal to {})", prefix, crit_under),
        (None, None) => unreachable!("unexpected got no over under"),
    }

    if matches.len() > 0 {
        println!("INFO: Matching processes:");
        for process in matches.iter().take(20) {
            println!("[{:>5}] {}", process.0, process.1.useful_cmdline());
        }
        if matches.len() > 20 {
            println!("And {} more...", matches.len() - 20)
        }
    }

    if args.kill_matching.is_some() {
        let signal = args.kill_matching.unwrap().0;
        let errors: Vec<_> = matches
            .iter()
            .map(|&(pid, _)| kill(*pid, signal))
            .filter(|result| result.is_err())
            .collect();
        if !errors.is_empty() {
            println!("INFO: There were {} errors killing processes", errors.len());
        }
    }
    if args.kill_matching_parents.is_some() {
        let signal = args.kill_matching_parents.unwrap().0;
        let mut parents = HashSet::new();
        let errors: Vec<_> = matches
            .iter()
            .map(|&(_, process)| process.stat.ppid)
            .filter(|ppid| parents.insert(*ppid))
            .map(|ppid| kill(ppid, signal))
            .filter(|result| result.is_err())
            .collect();
        if !errors.is_empty() {
            println!("INFO: There were {} errors killing processes", errors.len());
        }
    }

    status.exit();
}

/// Load currently running procs, and die if there is a surprising error
///
/// Normally if this can load *any* processes it returns what it can find, and
/// prints errors for procs that can't be parsed. But if `die_on_any_errors` is
/// true it dies if it cannot parse a *single* process.
fn load_procs(die_on_any_errors: bool) -> RunningProcs {
    match RunningProcs::currently_running() {
        Ok(procs) => procs,
        Err(ProcFsError::LoadProcsError(LoadProcsError { procs, errors })) => {
            let mut saw_real_error = false;
            for err in errors {
                match err {
                    // If the process no longer exists, that's fine
                    ProcFsError::Io(_) => {}
                    err => {
                        saw_real_error = true;
                        println!("WARN: Unexpected error loading some processes: {}", err);
                    }
                }
            }
            if die_on_any_errors && saw_real_error {
                Status::Critical.exit();
            }
            procs
        }
        Err(err) => {
            println!("ERROR: unable to load processes: {}", err);
            Status::Critical.exit();
        }
    }
}

/// Filter to processes that match the condition
///
/// The condition is:
///
/// * Does not match me or my parent process
/// * Does match the regex, if it is present
/// * Does match *any* of the states
fn filter_procs<'m>(
    re: &Option<Regex>,
    states: &[State],
    procs: &'m ProcMap,
) -> Vec<(&'m Pid, &'m Process)> {
    let me = getpid();
    let parent = getppid();

    debug!("filtering for procs that match: {:?} {:?}", re, states);

    procs
        .iter()
        .filter(|&(pid, _process)| *pid != me && *pid != parent)
        .filter(|&(_pid, process)| states.is_empty() || states.contains(&process.stat.state))
        .filter(|&(_pid, process)| {
            re.as_ref()
                .map(|re| {
                    let cmdline = process.useful_cmdline();
                    let is_match = re.is_match(&cmdline);
                    trace!("is_match={} process={}", is_match, cmdline);
                    is_match
                })
                .unwrap_or(true)
        })
        .collect()
}

#[cfg(test)]
mod unit {
    use structopt::StructOpt;

    use tabin_plugins::procfs::pid::Process;

    use super::*;

    #[test]
    fn validate_argparse() {
        let args = Args::from_iter(["c-p", "some.*proc", "--crit-under=1"].into_iter());
        assert_eq!(args.crit_under, Some(1));
    }

    #[test]
    fn validate_parse_zombies() {
        let args = Args::from_iter(["c-p", "some.*proc", "--state=zombie"].into_iter());
        assert_eq!(args.states, [State::Zombie]);
        let args =
            Args::from_iter(["c-p", "some.*proc", "--state=zombie", "--state", "S"].into_iter());
        assert_eq!(args.states, [State::Zombie, State::Sleeping]);
    }

    #[test]
    fn validate_parse_zombies_and_pattern() {
        let args = Args::from_iter(["c-p", "--state", "zombie", "--", "some.*proc"].into_iter());
        assert_eq!(args.states, [State::Zombie]);
        let args = Args::from_iter(["c-p", "--state=zombie", "some.*proc"].into_iter());
        assert_eq!(args.states, [State::Zombie]);
        let args = Args::from_iter(["c-p", "so.*proc", "--state", "zombie", "waiting"].into_iter());
        assert_eq!(args.states, [State::Zombie, State::Waiting]);
    }

    // Waiting for structopt 0.2.8 to be released with the from_iter_safe method
    // #[test]
    // #[should_panic]
    // fn regexes_must_be_valid() {
    //     Args::from_iter_safe(["c-p", "--crit-over=5", "["].into_iter()).unwrap();
    // }

    #[test]
    fn filter_procs_handles_patterns() {
        let mut procs = vec![Process::default(); 5];
        procs[2]
            .cmdline
            .raw
            .append(&mut vec!["hello".into(), "jar".into()]);
        let proc_map = vec_to_procmap(procs);
        let filtered = filter_procs(&regex("llo.*ar"), &[], &proc_map);
        assert_eq!(filtered.len(), 1);
    }

    #[test]
    fn filter_procs_handles_single_state() {
        let mut procs = vec![Process::default(); 5];
        procs[2].stat.state = State::Zombie;
        let proc_map = vec_to_procmap(procs);
        let filtered = filter_procs(&None, &[State::Zombie], &proc_map);
        assert_eq!(filtered.len(), 1);
    }

    #[test]
    fn filter_procs_handles_multiple_states() {
        let mut procs = vec![Process::default(); 5];
        procs[2].stat.state = State::Zombie;
        procs[3].stat.state = State::Stopped;
        let proc_map = vec_to_procmap(procs);

        let filtered = filter_procs(&None, &[State::Zombie], &proc_map);
        assert_eq!(filtered.len(), 1);

        let filtered = filter_procs(&None, &[State::Zombie, State::Stopped], &proc_map);
        assert_eq!(filtered.len(), 2);
    }

    #[test]
    fn filter_procs_handles_multiple_states_and_regex() {
        let mut procs = vec![Process::default(); 5];
        procs[2].stat.state = State::Zombie;
        procs[2].cmdline.raw.push("example".into());

        procs[3].stat.state = State::Stopped;
        let proc_map = vec_to_procmap(procs.clone());

        let filtered = filter_procs(&regex("exa"), &[], &proc_map);
        assert_eq!(filtered.len(), 1);

        let filtered = filter_procs(&regex("exa"), &[State::Zombie, State::Stopped], &proc_map);
        assert_eq!(filtered.len(), 1);
        assert_eq!(filtered[0].1, &procs[2]);

        let filtered = filter_procs(&regex("exa"), &[State::Stopped], &proc_map);
        assert_eq!(filtered.len(), 0);
    }

    fn regex(re: &str) -> Option<Regex> {
        Some(Regex::new(re).unwrap())
    }

    fn vec_to_procmap(procs: Vec<Process>) -> ProcMap {
        procs
            .into_iter()
            .enumerate()
            .map(|(i, process)| (Pid::from_raw(i as i32), process))
            .collect()
    }
}
<|endoftext|>


<commit_before>
extern crate rustc_serialize;
extern crate amqp;

use amqp::{ConsumerCallBackFn, Session, Table, Basic, Channel, Options};
use amqp::protocol;
use std::default::Default;
use rustc_serialize::json;
use std::slice;

// for debugging
fn listen_discovery(channel: &mut Channel) {
    fn participant_discovered(channel: &mut Channel, deliver: protocol::basic::Deliver,
                            headers: protocol::basic::BasicProperties, body: Vec<u8>){
        let s = std::str::from_utf8(&body).unwrap();
        let info: ParticipantInfo = json::decode(s).unwrap();
        println!("MsgFlo participant discovered: {:?}", info);
        //println!("Deliver info: {:?}", deliver);
        //println!("Content headers: {:?}", headers);
        //println!("Content body: ", );
        channel.basic_ack(deliver.delivery_tag, false);
    }

    let queue_name = "fbp";
    let consumer_name = channel.basic_consume(participant_discovered as ConsumerCallBackFn, queue_name,
                                                "", false, false, false, false, Table::new());
}

#[derive(Debug, Default, RustcDecodable, RustcEncodable)]
struct ParticipantPort {
    id: String, // port name
    queue: String, // the associated message queue
    // FIXME: support. Is a keyword so needs some special handling   type: String, // datatype, ex: "boolean"
    // options: queue options as specified by the message queue implementation
}

#[derive(Debug, Default, RustcDecodable, RustcEncodable)]
struct ParticipantInfo {
    id: String, // unique name
    role: String ,// role participant has
    component: String, // component the participant is instance of
//   label: Option<String>, // (optional) short human-readable description
//    icon: Option<String>, // (optional)
    inports: Vec<ParticipantPort>,
    outports: Vec<ParticipantPort>,
}

type ProcessFunction = fn(i32) -> i32;
struct Participant {
    info: ParticipantInfo,
    process: ProcessFunction,
}

struct Connection {
    session: Session,
    channel: Channel,
}

fn send_discovery(channel: &mut Channel, info: &ParticipantInfo) {
    let queue_name = "fbp"; // TODO: use an exchange istead, requires protocol change in msgflo

    let queue_declare = channel.queue_declare(queue_name, false, true, false, false, false, Table::new());
    let content_type = Some("application/json".to_string());
    let props = protocol::basic::BasicProperties { content_type: content_type, ..Default::default() };

    let payload = json::encode(&info).unwrap();
    println!("sending: {}", payload);
    let res = channel.basic_publish("", queue_name, true, false, props, payload.into_bytes());
}

fn start_participant(participant: &Participant) -> Connection {
    let mut session = Session::new(Options{vhost: "/", .. Default::default()}).ok().expect("Can't create session");
    let mut channel = session.open_channel(1).expect("channel");

    send_discovery(&mut channel, &participant.info);
    listen_discovery(&mut channel);

    channel.start_consuming();
    return Connection { session: session, channel: channel } ;
}

// for debugging
fn stop_participant(participant: &Participant, connection: &mut Connection) {

    let closed = connection.channel.close(200, "Bye".to_string());
    connection.session.close(200, "Good Bye".to_string());
}

fn main() {
    let info = ParticipantInfo {
        id: "part11".to_string(),
        role: "myrole".to_string(),
        component: "rust/First".to_string(),
        inports: Vec::<ParticipantPort>::new(),
        outports: Vec::<ParticipantPort>::new(),
    };
    fn process_repeat(input: i32) -> i32 {
        return input;
    }

    let p = Participant { info: info, process: process_repeat };

    let mut c = start_participant(&p);
    stop_participant(&p, &mut c);
}
<commit_msg>Almost ready to pass data through process()<commit_after>
extern crate rustc_serialize;
extern crate amqp;

use amqp::{ConsumerCallBackFn, Session, Table, Basic, Channel, Options};
use amqp::protocol;
use std::default::Default;
use rustc_serialize::json;
use std::slice;

// for debugging
fn listen_discovery(channel: &mut Channel) {
    fn participant_discovered(channel: &mut Channel, deliver: protocol::basic::Deliver,
                            headers: protocol::basic::BasicProperties, body: Vec<u8>){
        let s = std::str::from_utf8(&body).unwrap();
        let info: ParticipantInfo = json::decode(s).unwrap();
        println!("MsgFlo participant discovered: {:?}", info);
        //println!("Deliver info: {:?}", deliver);
        //println!("Content headers: {:?}", headers);
        //println!("Content body: ", );
        channel.basic_ack(deliver.delivery_tag, false);
    }

    let queue_name = "fbp";
    let consumer_name = channel.basic_consume(participant_discovered as ConsumerCallBackFn, queue_name,
                                                "", false, false, false, false, Table::new());
    println!("listening for discovery messages");
}

#[derive(Debug, Default, RustcDecodable, RustcEncodable)]
struct ParticipantPort {
    id: String, // port name
    queue: String, // the associated message queue
    // FIXME: support. Is a keyword so needs some special handling   type: String, // datatype, ex: "boolean"
    // options: queue options as specified by the message queue implementation
    // description?? TODO: standardize
}

#[derive(Debug, Default, RustcDecodable, RustcEncodable)]
struct ParticipantInfo {
    id: String, // unique name
    role: String ,// role participant has
    component: String, // component the participant is instance of
//   label: Option<String>, // (optional) short human-readable description
//    icon: Option<String>, // (optional)
    inports: Vec<ParticipantPort>,
    outports: Vec<ParticipantPort>,
}

type ProcessFunction = fn(i32) -> i32;
struct Participant {
    info: ParticipantInfo,
    process: ProcessFunction,
}

struct Connection {
    session: Session,
    channel: Channel,
}

// FIXME: wrap in protocol/command/payload envelope
fn create_queue_and_send(channel: &mut Channel, queue_name: &str, payload: String) {

    let queue_declare = channel.queue_declare(queue_name, false, true, false, false, false, Table::new());
    queue_declare.expect("queue creation failed");
    let content_type = Some("application/json".to_string());
    let props = protocol::basic::BasicProperties { content_type: content_type, ..Default::default() };
    println!("sending on {}: {}", queue_name, payload);
    let res = channel.basic_publish("", queue_name, true, false, props, payload.into_bytes());
    res.expect("send on new queue failed");
}

fn send_discovery(channel: &mut Channel, info: &ParticipantInfo) {
    let queue_name = "fbp"; // TODO: use an exchange istead, requires protocol change in msgflo

    let payload = json::encode(&info).unwrap();
    create_queue_and_send(channel, queue_name, payload);
}

// FIXME: actually call ProcessFunction
fn setup_inport(port: &ParticipantPort, connection: &mut Connection) {
    println!("setup inport: {}", port.queue.to_string());

    fn on_message(channel: &mut Channel, deliver: protocol::basic::Deliver,
                            headers: protocol::basic::BasicProperties, body: Vec<u8>) {

        println!("!! got msg on inport");
        let s = std::str::from_utf8(&body).unwrap();
        let info: ParticipantInfo = json::decode(s).unwrap();
        println!("got message on: {:?}", info);
        channel.basic_ack(deliver.delivery_tag, false);
    }

    // create
    let declare = connection.channel.queue_declare(port.queue.to_string(), false, true, false, false, false, Table::new());
    declare.expect("inport queue creation failed");

    // subscribe
    let q = port.queue.to_string();
    let cons = connection.channel.basic_consume(on_message as ConsumerCallBackFn, q,
                                                "".to_string(), false, false, false, false, Table::new());

    cons.expect("iport setup failed");

    println!("setup inport done: {:?}", port.id.to_string());
}

fn setup_outport(port: &ParticipantPort, connection: &mut Connection) {

    let queue_declare = connection.channel.queue_declare(port.queue.to_string(), false, true, false, false, false, Table::new());
    let content_type = Some("application/json".to_string());
}

// FIXME: actually send
fn send_output(participant: &Participant) {

    let queue_name = participant.info.outports[0].queue.to_string();
//    let payload = json::encode(&info).unwrap();
//    println!("sending: {}", payload);
//    let props = protocol::basic::BasicProperties { content_type: content_type, ..Default::default() };
//    let res = channel.basic_publish("", port.queue.to_string(), true, false, props, payload.into_bytes());
}

fn start_participant(participant: &Participant) -> Connection {
    let mut session = Session::new(Options{vhost: "/", .. Default::default()}).ok().expect("Can't create session");
    let mut channel = session.open_channel(1).expect("channel");

    send_discovery(&mut channel, &participant.info);

    let mut conn = Connection { session: session, channel: channel };

    // XXX: seems like can only connect to one queue??
    listen_discovery(&mut conn.channel); // TESTING

    setup_inport(&participant.info.inports[0], &mut conn);

    return conn;
}

// for debugging
fn stop_participant(participant: &Participant, connection: &mut Connection) {

    let closed = connection.channel.close(200, "Bye".to_string());
    connection.session.close(200, "Good Bye".to_string());
}

// TODO: wire up input/output to the process() function
// TODO: respect MSGFLO_BROKER envvar
// TODO: setup the msgflo hetro automated test
// TODO: nicer way to declare ports
fn main() {
    let info = ParticipantInfo {
        id: "part11".to_string(),
        role: "myrole".to_string(),
        component: "rust/First".to_string(),
        inports: vec! [ ParticipantPort { id: "in".to_string(), queue: "rustparty.IN".to_string() } ],
        outports: vec! [ ParticipantPort { id: "out".to_string(), queue: "rustparty.OUT".to_string() } ],
    };
    fn process_repeat(input: i32) -> i32 {
        return input;
    }

    let p = Participant { info: info, process: process_repeat };

    let mut c = start_participant(&p);

    create_queue_and_send(&mut c.channel, "rustparty.in", "{ \"data\": 300 }".to_string());

    c.channel.start_consuming();
    stop_participant(&p, &mut c);
}
<|endoftext|>


<commit_before>#![deny(warnings)]

extern crate raw_cpuid;
extern crate syscall;

use std::env;
use std::fs::File;
use std::io::{self, Read, Write};

fn main() {
    let user = env::var("USER").unwrap_or(String::new());
    let mut hostname = String::new();
    if let Ok(mut file) = File::open("/etc/hostname") {
        let _ = file.read_to_string(&mut hostname);
    }

    let mut uptime = 0;

    let mut ts = syscall::TimeSpec::default();
    if syscall::clock_gettime(syscall::CLOCK_MONOTONIC, &mut ts).is_ok() {
        uptime = ts.tv_sec;
    }

    let mut width = 0;
    let mut height = 0;
    if let Ok(display_name) = env::var("DISPLAY") {
        if let Ok(display) = syscall::open(&display_name, syscall::O_STAT) {
            let mut buf: [u8; 4096] = [0; 4096];
            if let Ok(count) = syscall::fpath(display, &mut buf) {
                let path = unsafe { String::from_utf8_unchecked(Vec::from(&buf[..count])) };
                let res = path.split(":").nth(1).unwrap_or("");
                width = res.split("/").nth(1).unwrap_or("").parse::<i32>().unwrap_or(0);
                height = res.split("/").nth(2).unwrap_or("").parse::<i32>().unwrap_or(0);
            }
            let _ = syscall::close(display);
        }
    }

    let mut cpu = String::new();
    {
        let cpuid = raw_cpuid::CpuId::new();
        if let Some(info) = cpuid.get_extended_function_info() {
            if let Some(brand) = info.processor_brand_string() {
                cpu = brand.to_string();
            }
        }
    }

    let mut ram = String::new();
    {
        let mut stat = syscall::StatVfs::default();
        if let Ok(fd) = syscall::open("memory:", syscall::O_STAT) {
            if syscall::fstatvfs(fd, &mut stat).is_ok() {
                let size = stat.f_blocks * stat.f_bsize as u64;
                let used = (stat.f_blocks - stat.f_bfree) * stat.f_bsize as u64;

                ram = format!("{}MB / {}MB", (used + 1048575)/1048576, (size + 1048575)/1048576);
            }
            let _ = syscall::close(fd);
        }
    }

    let left = [
        "              :+yMMMMy+:              ",
        "         .+dddNMMMMMMMMNddd+.         ",
        "       sydMMMMo/sMMMMs/oMMMMdys       ",
        "    .oMMMdso              osdMMMo.    ",
        "   .+MMd/`   -:::::::`      `/dMM+.   ",
        "  +dMMN.     NMMNNNMMdyo`     .NMMd+  ",
        "  yMMN       NM+   oomMMN       NMMy  ",
        " hNMd.       NM+     `oMN       .dMNh ",
        " dMMh        NM+     `oMN        hMMd ",
        " dMMh        NM+-oooodMMN        hMMd ",
        " dMMh        NM+/MMMMdhs`        hMMd ",
        " hNMd.       NM+`/mMMm+         .dMNh ",
        "  yMMN       NM+   oNMd+        NMMy  ",
        "  +dMMN.     NM+    omMMm     .NMMd+  ",
        "   .+MMd/`   --.     .---   `/dMM+.   ",
        "    .oMMMdso              osdMMMo.    ",
        "       sydMMMMo////////oMMMMdys       ",
        "         .+dddNMMMMMMMMNddd+.         ",
        "              :++++++++:              ",
    ];

    let right = [
        format!("\x1B[1;38;5;75m{}\x1B[0m@\x1B[1;38;5;75m{}\x1B[0m", user, hostname.trim()),
        format!("\x1B[1;38;5;75mOS:\x1B[0m redox-os"),
        format!("\x1B[1;38;5;75mKernel:\x1B[0m redox"),
        format!("\x1B[1;38;5;75mUptime:\x1B[0m {}s", uptime),
        format!("\x1B[1;38;5;75mShell:\x1B[0m ion"),
        format!("\x1B[1;38;5;75mResolution:\x1B[0m {}x{}", width, height),
        format!("\x1B[1;38;5;75mDE:\x1B[0m orbital"),
        format!("\x1B[1;38;5;75mWM:\x1B[0m orbital"),
        format!("\x1B[1;38;5;75mFont:\x1B[0m unifont"),
        format!("\x1B[1;38;5;75mCPU:\x1B[0m {}", cpu),
        format!("\x1B[1;38;5;75mRAM:\x1B[0m {}", ram)
    ];

    let mut string = String::new();
    for i in 0..left.len() {
        string.push_str("\x1B[1;38;5;75m");
        string.push_str(left[i]);
        string.push_str("  \x1B[0m");
        if let Some(r) = right.get(i) {
            string.push_str(r);
        }
        string.push_str("\n");
    }

    io::stdout().write(string.as_bytes()).unwrap();
}
<commit_msg>screenfetch: Trim trailing whitespace from CPU The raw_cpuid crate passes the CPU brand string with trailing whitespace which borks the output. According to the tests of that crate it's an expected behaviour so the whitespace needs to be removed here.<commit_after>#![deny(warnings)]

extern crate raw_cpuid;
extern crate syscall;

use std::env;
use std::fs::File;
use std::io::{self, Read, Write};

fn main() {
    let user = env::var("USER").unwrap_or(String::new());
    let mut hostname = String::new();
    if let Ok(mut file) = File::open("/etc/hostname") {
        let _ = file.read_to_string(&mut hostname);
    }

    let mut uptime = 0;

    let mut ts = syscall::TimeSpec::default();
    if syscall::clock_gettime(syscall::CLOCK_MONOTONIC, &mut ts).is_ok() {
        uptime = ts.tv_sec;
    }

    let mut width = 0;
    let mut height = 0;
    if let Ok(display_name) = env::var("DISPLAY") {
        if let Ok(display) = syscall::open(&display_name, syscall::O_STAT) {
            let mut buf: [u8; 4096] = [0; 4096];
            if let Ok(count) = syscall::fpath(display, &mut buf) {
                let path = unsafe { String::from_utf8_unchecked(Vec::from(&buf[..count])) };
                let res = path.split(":").nth(1).unwrap_or("");
                width = res.split("/").nth(1).unwrap_or("").parse::<i32>().unwrap_or(0);
                height = res.split("/").nth(2).unwrap_or("").parse::<i32>().unwrap_or(0);
            }
            let _ = syscall::close(display);
        }
    }

    let mut cpu = String::new();
    {
        let cpuid = raw_cpuid::CpuId::new();
        if let Some(info) = cpuid.get_extended_function_info() {
            if let Some(brand) = info.processor_brand_string() {
                cpu = brand.trim().to_string();
            }
        }
    }

    let mut ram = String::new();
    {
        let mut stat = syscall::StatVfs::default();
        if let Ok(fd) = syscall::open("memory:", syscall::O_STAT) {
            if syscall::fstatvfs(fd, &mut stat).is_ok() {
                let size = stat.f_blocks * stat.f_bsize as u64;
                let used = (stat.f_blocks - stat.f_bfree) * stat.f_bsize as u64;

                ram = format!("{}MB / {}MB", (used + 1048575)/1048576, (size + 1048575)/1048576);
            }
            let _ = syscall::close(fd);
        }
    }

    let left = [
        "              :+yMMMMy+:              ",
        "         .+dddNMMMMMMMMNddd+.         ",
        "       sydMMMMo/sMMMMs/oMMMMdys       ",
        "    .oMMMdso              osdMMMo.    ",
        "   .+MMd/`   -:::::::`      `/dMM+.   ",
        "  +dMMN.     NMMNNNMMdyo`     .NMMd+  ",
        "  yMMN       NM+   oomMMN       NMMy  ",
        " hNMd.       NM+     `oMN       .dMNh ",
        " dMMh        NM+     `oMN        hMMd ",
        " dMMh        NM+-oooodMMN        hMMd ",
        " dMMh        NM+/MMMMdhs`        hMMd ",
        " hNMd.       NM+`/mMMm+         .dMNh ",
        "  yMMN       NM+   oNMd+        NMMy  ",
        "  +dMMN.     NM+    omMMm     .NMMd+  ",
        "   .+MMd/`   --.     .---   `/dMM+.   ",
        "    .oMMMdso              osdMMMo.    ",
        "       sydMMMMo////////oMMMMdys       ",
        "         .+dddNMMMMMMMMNddd+.         ",
        "              :++++++++:              ",
    ];

    let right = [
        format!("\x1B[1;38;5;75m{}\x1B[0m@\x1B[1;38;5;75m{}\x1B[0m", user, hostname.trim()),
        format!("\x1B[1;38;5;75mOS:\x1B[0m redox-os"),
        format!("\x1B[1;38;5;75mKernel:\x1B[0m redox"),
        format!("\x1B[1;38;5;75mUptime:\x1B[0m {}s", uptime),
        format!("\x1B[1;38;5;75mShell:\x1B[0m ion"),
        format!("\x1B[1;38;5;75mResolution:\x1B[0m {}x{}", width, height),
        format!("\x1B[1;38;5;75mDE:\x1B[0m orbital"),
        format!("\x1B[1;38;5;75mWM:\x1B[0m orbital"),
        format!("\x1B[1;38;5;75mFont:\x1B[0m unifont"),
        format!("\x1B[1;38;5;75mCPU:\x1B[0m {}", cpu),
        format!("\x1B[1;38;5;75mRAM:\x1B[0m {}", ram)
    ];

    let mut string = String::new();
    for i in 0..left.len() {
        string.push_str("\x1B[1;38;5;75m");
        string.push_str(left[i]);
        string.push_str("  \x1B[0m");
        if let Some(r) = right.get(i) {
            string.push_str(r);
        }
        string.push_str("\n");
    }

    io::stdout().write(string.as_bytes()).unwrap();
}
<|endoftext|>


<commit_before>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io

extern crate rustc_serialize;
extern crate docopt;
extern crate csv;
extern crate mimir;
extern crate itertools;
extern crate mimirsbrunn;
#[macro_use]
extern crate log;

use std::rc::Rc;
use mimir::rubber::Rubber;
use docopt::Docopt;
use mimirsbrunn::utils::{format_label, get_zip_codes_from_admins};
use mimirsbrunn::admin_geofinder::AdminGeoFinder;
use std::collections::HashMap;

const USAGE: &'static str = "
Usage:
    stops2mimir --help
    stops2mimir --input=<file> \
     [--connection-string=<connection-string>] [--dataset=<dataset>] [--city-level=<level>]

Options:
    -h, --help                Show this message.
    -i, --input=<file>        NTFS stops.txt file.
    -c, --connection-string=<connection-string>   \
                              Elasticsearch parameters [default: http://localhost:9200/munin].
    -d, --dataset=<dataset>   Name of the dataset [default: fr].
    -C, --city-level=<level>  City level to calculate weight [default: 8]
";

#[derive(Debug, RustcDecodable)]
struct Args {
    flag_input: String,
    flag_dataset: String,
    flag_connection_string: String,
    flag_city_level: u32,
}

struct StopPointIter<'a, R: std::io::Read + 'a> {
    iter: csv::StringRecords<'a, R>,
    stop_id_pos: usize,
    stop_lat_pos: usize,
    stop_lon_pos: usize,
    stop_name_pos: usize,
    location_type_pos: Option<usize>,
    stop_visible_pos: Option<usize>,
    parent_station_pos: Option<usize>,
    nb_stop_points: &'a mut HashMap<String, u32>,
}

impl<'a, R: std::io::Read + 'a> StopPointIter<'a, R> {
    fn new(r: &'a mut csv::Reader<R>,
           nb_stop_points: &'a mut HashMap<String, u32>)
           -> csv::Result<Self> {
        let headers = try!(r.headers());
        let get_optional_pos = |name| headers.iter().position(|s| s == name);

        let get_pos = |field| {
            get_optional_pos(field).ok_or_else(|| {
                csv::Error::Decode(format!("Invalid file, cannot find column '{}'", field))
            })
        };

        Ok(StopPointIter {
            iter: r.records(),
            stop_id_pos: try!(get_pos("stop_id")),
            stop_lat_pos: try!(get_pos("stop_lat")),
            stop_lon_pos: try!(get_pos("stop_lon")),
            stop_name_pos: try!(get_pos("stop_name")),
            location_type_pos: get_optional_pos("location_type"),
            stop_visible_pos: get_optional_pos("visible"),
            parent_station_pos: get_optional_pos("parent_station"),
            nb_stop_points: nb_stop_points,
        })
    }
    fn get_location_type(&self, record: &[String]) -> Option<u8> {
        self.location_type_pos.and_then(|pos| record.get(pos).and_then(|s| s.parse().ok()))
    }
    fn get_visible(&self, record: &[String]) -> Option<u8> {
        self.stop_visible_pos.and_then(|pos| record.get(pos).and_then(|s| s.parse().ok()))
    }

    fn get_parent_station<'b>(&self, record: &'b [String]) -> Option<&'b String> {
        self.parent_station_pos.and_then(|pos| record.get(pos))
    }
}

impl<'a, R: std::io::Read + 'a> Iterator for StopPointIter<'a, R> {
    type Item = csv::Result<mimir::Stop>;
    fn next(&mut self) -> Option<Self::Item> {
        fn get(record: &[String], pos: usize) -> csv::Result<&str> {
            record.get(pos)
                .map(|s| s.as_str())
                .ok_or_else(|| csv::Error::Decode(format!("Failed accessing record '{}'.", pos)))
        }
        fn parse_f64(s: &str) -> csv::Result<f64> {
            s.parse()
                .map_err(|_| csv::Error::Decode(format!("Failed converting '{}' from str.", s)))
        }

        fn is_valid_stop_area(location_type: &Option<u8>, visible: &Option<u8>) -> csv::Result<()> {
            if *location_type != Some(1) {
                Err(csv::Error::Decode("not a stop_area.".to_string()))
            } else if *visible == Some(0) {
                Err(csv::Error::Decode("stop_area invisible.".to_string()))
            } else {
                Ok(())
            }
        }

        self.iter.next().map(|r| {
            r.and_then(|r| {
                let location_type = self.get_location_type(&r);
                let visible = self.get_visible(&r);
                let parent_station = self.get_parent_station(&r);
                //if it's a stop point, we update its stop_area counter
                if let (Some(0), Some(id)) = (location_type, parent_station) {
                    if !id.is_empty() {
                        *self.nb_stop_points.entry(format!("stop_area:{}", id)).or_insert(0) += 1;
                    }
                }
                try!(is_valid_stop_area(&location_type, &visible));
                let stop_id = try!(get(&r, self.stop_id_pos));
                let stop_lat = try!(get(&r, self.stop_lat_pos));
                let stop_lat = try!(parse_f64(stop_lat));
                let stop_lon = try!(get(&r, self.stop_lon_pos));
                let stop_lon = try!(parse_f64(stop_lon));
                let stop_name = try!(get(&r, self.stop_name_pos));
                Ok(mimir::Stop {
                    id: format!("stop_area:{}", stop_id), // prefix to match navitia's id
                    coord: mimir::Coord::new(stop_lat, stop_lon),
                    label: stop_name.to_string(),
                    weight: 1,
                    zip_codes: vec![],
                    administrative_regions: vec![],
                    name: stop_name.to_string(),
                })
            })
        })
    }
}

fn attach_stop(stop: &mut mimir::Stop, admins: Vec<Rc<mimir::Admin>>, city_level: u32) {
    stop.administrative_regions = admins;
    stop.label = format_label(&stop.administrative_regions, city_level, &stop.name);
    stop.zip_codes = get_zip_codes_from_admins(&stop.administrative_regions);
}

/// Attach the stops to administrative regions
///
/// The admins are loaded from Elasticsearch and stored in a quadtree
/// We attach a stop with all the admins that have a boundary containing
/// the coordinate of the stop
fn attach_stops_to_admins<'a, It: Iterator<Item = &'a mut mimir::Stop>>(stops: It,
                                                                        rubber: &mut Rubber,
                                                                        city_level: u32) {
    let admins = rubber.get_all_admins().unwrap_or_else(|_| {
        info!("Administratives regions not found in elasticsearch db");
        vec![]
    });

    info!("{} administrative regions loaded from mimir", admins.len());

    let admins_geofinder = admins.into_iter().collect::<AdminGeoFinder>();

    let mut nb_unmatched = 0u32;
    let mut nb_matched = 0u32;
    for mut stop in stops {
        let admins = admins_geofinder.get(&stop.coord);

        if admins.is_empty() {
            nb_unmatched += 1;
        } else {
            nb_matched += 1;
        }

        attach_stop(&mut stop, admins, city_level);
    }

    info!("there is {}/{} stops without any admin",
          nb_unmatched,
          nb_matched + nb_unmatched);
}

// Update weight value for each stop_area from HashMap.
fn finalize_stop_area_weight<'a, It: Iterator<Item = &'a mut mimir::Stop>>(stops: It,
    nb_stop_points: &HashMap<String, u32>) {
    for stop in stops {
        if let Some(weight) = nb_stop_points.get(&stop.id) {
            stop.weight = *weight;
        }
    }
}

fn main() {
    mimir::logger_init().unwrap();
    info!("Launching stops2mimir...");

    let args: Args = Docopt::new(USAGE).and_then(|dopt| dopt.decode()).unwrap_or_else(|e| e.exit());

    info!("creation of indexes");
    let mut rubber = Rubber::new(&args.flag_connection_string);
    let mut rdr = csv::Reader::from_file(args.flag_input).unwrap().double_quote(true);

    let mut nb_stop_points = HashMap::new();

    let mut stops: Vec<mimir::Stop> = StopPointIter::new(&mut rdr, &mut nb_stop_points)
        .unwrap()
        .filter_map(|rc| {
            rc.map_err(|e| debug!("skip csv line because: {}", e))
                .ok()
        })
        .collect();

    attach_stops_to_admins(stops.iter_mut(), &mut rubber, args.flag_city_level);

    finalize_stop_area_weight(stops.iter_mut(), &nb_stop_points);

    info!("Importing stops into Mimir");
    let nb_stops = rubber.index(&args.flag_dataset, stops.iter()).unwrap();

    info!("Nb of indexed stops: {}", nb_stops);

}

#[test]
fn test_load_stops() {
    use itertools::Itertools;
    // stops.txt:
    // SP:main_station : StopPoint object
    // SA:main_station: StopArea valid
    // SA:second_station: StopArea valid with visible is empty
    // SA:invisible_station: invisible StopArea
    // SA:without_lat: StopArea object without lattitude coord
    // SA:witout_lon: StopArea object without longitude coord
    // SA:station_no_city: StopArea far away, we won't be able to attach it to a city
    let mut rdr = csv::Reader::from_file("./tests/fixtures/stops.txt".to_string())
        .unwrap()
        .double_quote(true);

    let mut nb_stop_points = HashMap::new();
    let stops: Vec<mimir::Stop> = StopPointIter::new(&mut rdr, &mut nb_stop_points)
        .unwrap()
        .filter_map(|rc| {
            rc.map_err(|e| println!("error at csv line decoding : {}", e))
                .ok()
        })
        .collect();
    assert_eq!(stops.len(), 5);
    let ids: Vec<_> = stops.iter().map(|s| s.id.clone()).sorted();
    assert_eq!(ids,
               vec!["stop_area:SA:main_station",
                    "stop_area:SA:second_station",
                    "stop_area:SA:station_no_city",
                    "stop_area:SA:weight_1_station",
                    "stop_area:SA:weight_3_station"]);
}
<commit_msg>mimir: fromat<commit_after>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io

extern crate rustc_serialize;
extern crate docopt;
extern crate csv;
extern crate mimir;
extern crate itertools;
extern crate mimirsbrunn;
#[macro_use]
extern crate log;

use std::rc::Rc;
use mimir::rubber::Rubber;
use docopt::Docopt;
use mimirsbrunn::utils::{format_label, get_zip_codes_from_admins};
use mimirsbrunn::admin_geofinder::AdminGeoFinder;
use std::collections::HashMap;

const USAGE: &'static str = "
Usage:
    stops2mimir --help
    stops2mimir --input=<file> \
     [--connection-string=<connection-string>] [--dataset=<dataset>] [--city-level=<level>]

Options:
    -h, --help                Show this message.
    -i, --input=<file>        NTFS stops.txt file.
    -c, --connection-string=<connection-string>   \
                              Elasticsearch parameters [default: http://localhost:9200/munin].
    -d, --dataset=<dataset>   Name of the dataset [default: fr].
    -C, --city-level=<level>  City level to calculate weight [default: 8]
";

#[derive(Debug, RustcDecodable)]
struct Args {
    flag_input: String,
    flag_dataset: String,
    flag_connection_string: String,
    flag_city_level: u32,
}

struct StopPointIter<'a, R: std::io::Read + 'a> {
    iter: csv::StringRecords<'a, R>,
    stop_id_pos: usize,
    stop_lat_pos: usize,
    stop_lon_pos: usize,
    stop_name_pos: usize,
    location_type_pos: Option<usize>,
    stop_visible_pos: Option<usize>,
    parent_station_pos: Option<usize>,
    nb_stop_points: &'a mut HashMap<String, u32>,
}

impl<'a, R: std::io::Read + 'a> StopPointIter<'a, R> {
    fn new(r: &'a mut csv::Reader<R>,
           nb_stop_points: &'a mut HashMap<String, u32>)
           -> csv::Result<Self> {
        let headers = try!(r.headers());
        let get_optional_pos = |name| headers.iter().position(|s| s == name);

        let get_pos = |field| {
            get_optional_pos(field).ok_or_else(|| {
                csv::Error::Decode(format!("Invalid file, cannot find column '{}'", field))
            })
        };

        Ok(StopPointIter {
            iter: r.records(),
            stop_id_pos: try!(get_pos("stop_id")),
            stop_lat_pos: try!(get_pos("stop_lat")),
            stop_lon_pos: try!(get_pos("stop_lon")),
            stop_name_pos: try!(get_pos("stop_name")),
            location_type_pos: get_optional_pos("location_type"),
            stop_visible_pos: get_optional_pos("visible"),
            parent_station_pos: get_optional_pos("parent_station"),
            nb_stop_points: nb_stop_points,
        })
    }
    fn get_location_type(&self, record: &[String]) -> Option<u8> {
        self.location_type_pos.and_then(|pos| record.get(pos).and_then(|s| s.parse().ok()))
    }
    fn get_visible(&self, record: &[String]) -> Option<u8> {
        self.stop_visible_pos.and_then(|pos| record.get(pos).and_then(|s| s.parse().ok()))
    }

    fn get_parent_station<'b>(&self, record: &'b [String]) -> Option<&'b String> {
        self.parent_station_pos.and_then(|pos| record.get(pos))
    }
}

impl<'a, R: std::io::Read + 'a> Iterator for StopPointIter<'a, R> {
    type Item = csv::Result<mimir::Stop>;
    fn next(&mut self) -> Option<Self::Item> {
        fn get(record: &[String], pos: usize) -> csv::Result<&str> {
            record.get(pos)
                .map(|s| s.as_str())
                .ok_or_else(|| csv::Error::Decode(format!("Failed accessing record '{}'.", pos)))
        }
        fn parse_f64(s: &str) -> csv::Result<f64> {
            s.parse()
                .map_err(|_| csv::Error::Decode(format!("Failed converting '{}' from str.", s)))
        }

        fn is_valid_stop_area(location_type: &Option<u8>, visible: &Option<u8>) -> csv::Result<()> {
            if *location_type != Some(1) {
                Err(csv::Error::Decode("not a stop_area.".to_string()))
            } else if *visible == Some(0) {
                Err(csv::Error::Decode("stop_area invisible.".to_string()))
            } else {
                Ok(())
            }
        }

        self.iter.next().map(|r| {
            r.and_then(|r| {
                let location_type = self.get_location_type(&r);
                let visible = self.get_visible(&r);
                let parent_station = self.get_parent_station(&r);
                //if it's a stop point, we update its stop_area counter
                if let (Some(0), Some(id)) = (location_type, parent_station) {
                    if !id.is_empty() {
                        *self.nb_stop_points.entry(format!("stop_area:{}", id)).or_insert(0) += 1;
                    }
                }
                try!(is_valid_stop_area(&location_type, &visible));
                let stop_id = try!(get(&r, self.stop_id_pos));
                let stop_lat = try!(get(&r, self.stop_lat_pos));
                let stop_lat = try!(parse_f64(stop_lat));
                let stop_lon = try!(get(&r, self.stop_lon_pos));
                let stop_lon = try!(parse_f64(stop_lon));
                let stop_name = try!(get(&r, self.stop_name_pos));
                Ok(mimir::Stop {
                    id: format!("stop_area:{}", stop_id), // prefix to match navitia's id
                    coord: mimir::Coord::new(stop_lat, stop_lon),
                    label: stop_name.to_string(),
                    weight: 1,
                    zip_codes: vec![],
                    administrative_regions: vec![],
                    name: stop_name.to_string(),
                })
            })
        })
    }
}

fn attach_stop(stop: &mut mimir::Stop, admins: Vec<Rc<mimir::Admin>>, city_level: u32) {
    stop.administrative_regions = admins;
    stop.label = format_label(&stop.administrative_regions, city_level, &stop.name);
    stop.zip_codes = get_zip_codes_from_admins(&stop.administrative_regions);
}

/// Attach the stops to administrative regions
///
/// The admins are loaded from Elasticsearch and stored in a quadtree
/// We attach a stop with all the admins that have a boundary containing
/// the coordinate of the stop
fn attach_stops_to_admins<'a, It: Iterator<Item = &'a mut mimir::Stop>>(stops: It,
                                                                        rubber: &mut Rubber,
                                                                        city_level: u32) {
    let admins = rubber.get_all_admins().unwrap_or_else(|_| {
        info!("Administratives regions not found in elasticsearch db");
        vec![]
    });

    info!("{} administrative regions loaded from mimir", admins.len());

    let admins_geofinder = admins.into_iter().collect::<AdminGeoFinder>();

    let mut nb_unmatched = 0u32;
    let mut nb_matched = 0u32;
    for mut stop in stops {
        let admins = admins_geofinder.get(&stop.coord);

        if admins.is_empty() {
            nb_unmatched += 1;
        } else {
            nb_matched += 1;
        }

        attach_stop(&mut stop, admins, city_level);
    }

    info!("there is {}/{} stops without any admin",
          nb_unmatched,
          nb_matched + nb_unmatched);
}

// Update weight value for each stop_area from HashMap.
fn finalize_stop_area_weight<'a, It: Iterator<Item = &'a mut mimir::Stop>>(stops: It,
    nb_stop_points: &HashMap<String, u32>) {
    for stop in stops {
        if let Some(weight) = nb_stop_points.get(&stop.id) {
            stop.weight = *weight;
        }
    }
}

fn main() {
    mimir::logger_init().unwrap();
    info!("Launching stops2mimir...");

    let args: Args = Docopt::new(USAGE).and_then(|dopt| dopt.decode()).unwrap_or_else(|e| e.exit());

    info!("creation of indexes");
    let mut rubber = Rubber::new(&args.flag_connection_string);
    let mut rdr = csv::Reader::from_file(args.flag_input).unwrap().double_quote(true);

    let mut nb_stop_points = HashMap::new();

    let mut stops: Vec<mimir::Stop> = StopPointIter::new(&mut rdr, &mut nb_stop_points)
        .unwrap()
        .filter_map(|rc| {
            rc.map_err(|e| debug!("skip csv line because: {}", e))
                .ok()
        })
        .collect();

    attach_stops_to_admins(stops.iter_mut(), &mut rubber, args.flag_city_level);

    finalize_stop_area_weight(stops.iter_mut(), &nb_stop_points);

    info!("Importing stops into Mimir");
    let nb_stops = rubber.index(&args.flag_dataset, stops.iter()).unwrap();

    info!("Nb of indexed stops: {}", nb_stops);

}

#[test]
fn test_load_stops() {
    use itertools::Itertools;
    // stops.txt:
    // SP:main_station : StopPoint object
    // SA:main_station: StopArea valid
    // SA:second_station: StopArea valid with visible is empty
    // SA:invisible_station: invisible StopArea
    // SA:without_lat: StopArea object without lattitude coord
    // SA:witout_lon: StopArea object without longitude coord
    // SA:station_no_city: StopArea far away, we won't be able to attach it to a city
    let mut rdr = csv::Reader::from_file("./tests/fixtures/stops.txt".to_string())
        .unwrap()
        .double_quote(true);

    let mut nb_stop_points = HashMap::new();
    let stops: Vec<mimir::Stop> = StopPointIter::new(&mut rdr, &mut nb_stop_points)
        .unwrap()
        .filter_map(|rc| {
            rc.map_err(|e| println!("error at csv line decoding : {}", e))
                .ok()
        })
        .collect();
    assert_eq!(stops.len(), 5);
    let ids: Vec<_> = stops.iter().map(|s| s.id.clone()).sorted();
    assert_eq!(ids,
               vec!["stop_area:SA:main_station",
                    "stop_area:SA:second_station",
                    "stop_area:SA:station_no_city",
                    "stop_area:SA:weight_1_station",
                    "stop_area:SA:weight_3_station"]);
}
<|endoftext|>


<commit_before>// (C)opyleft 2013 Frank Denis

/*!
 * Bloom filter for Rust
 *
 * This is a simple but fast Bloom filter implementation, that requires only
 * 2 hash functions, generated with SipHash-2-4 using randomized keys.
 */

#[desc = "A fast Bloom filter implementation."];
#[license = "BSD"];
#[crate_id="bloomfilter#0.1"];
#[crate_type = "rlib"];
#[warn(non_camel_case_types, non_uppercase_statics, non_uppercase_statics, unnecessary_qualification, managed_heap_memory)];

extern crate collections;
extern crate extra;

use std::cmp;
use std::num;
use std::rand;
use std::rand::Rng;
use collections::bitv;

struct SipHashKey {
    k1: u64,
    k2: u64
}

impl SipHashKey {
    fn new_random() -> SipHashKey {
        let mut rng = rand::task_rng();
        SipHashKey {
            k1: rand::Rand::rand(& mut rng),
            k2: rand::Rand::rand(& mut rng)
        }
    }
}

/// Bloom filter structure
pub struct Bloom {
    priv bitmap: bitv::Bitv,
    priv bitmap_bits: u64,
    priv k_num: uint,
    priv skeys: [SipHashKey, ..2]
}

impl Bloom {

/// Create a new bloom filter structure.
/// bitmap_size is the size in bytes (not bits) that will be allocated in memory
/// items_count is an estimation of the maximum number of items to store.
    pub fn new(bitmap_size: uint, items_count: uint) -> Bloom {
        assert!(bitmap_size > 0u && items_count > 0u);
        let bitmap_bits = (bitmap_size as u64) * 8u64;
        let k_num = Bloom::optimal_k_num(bitmap_bits, items_count);
        let bitmap = bitv::Bitv::new(bitmap_bits as uint, false);
        let skeys = [ SipHashKey::new_random(), SipHashKey::new_random() ];
        Bloom {
            bitmap: bitmap,
            bitmap_bits: bitmap_bits,
            k_num: k_num,
            skeys: skeys
        }
    }

/// Create a new bloom filter structure.
/// items_count is an estimation of the maximum number of items to store.
/// fp_p is the wanted rate of false positives, in ]0.0, 1.0[
    pub fn new_for_fp_rate(items_count: uint, fp_p: f64) -> Bloom {
        let bitmap_size = Bloom::compute_bitmap_size(items_count, fp_p);
        Bloom::new(bitmap_size, items_count)
    }

/// Compute a recommended bitmap size for items_count items
/// and a fp_p rate of false positives.
/// fp_p obviously has to be within the ]0.0, 1.0[ range.
    pub fn compute_bitmap_size(items_count: uint, fp_p: f64) -> uint {
        assert!(items_count > 0u);
        assert!(fp_p > 0.0 && fp_p < 1.0);
        let log2 = num::ln(2.0f64);
        let log2_2 = log2 * log2;
        ((items_count as f64) * num::ln(fp_p) / (-8.0 * log2_2)).ceil() as uint
    }

/// Record the presence of an item.
    pub fn set<T: Hash>(& mut self, item: T) {
        let mut hashes = [ 0u64, 0u64 ];
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            self.bitmap.set(bit_offset, true);
        }
    }

/// Check if an item is present in the set.
/// There can be false positives, but no false negatives.
    pub fn check<T: Hash>(&self, item: T) -> bool {
        let mut hashes = [ 0u64, 0u64 ];
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            if self.bitmap.get(bit_offset) == false {
                return false;
            }
        }
        true
    }

/// Record the presence of an item in the set,
/// and return the previous state of this item.
    pub fn check_and_set<T: Hash>(&mut self, item: T) -> bool {
        let mut hashes = [ 0u64, 0u64 ];
        let mut found = true;
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            if self.bitmap.get(bit_offset) == false {
                found = false;
                self.bitmap.set(bit_offset, true);
            }
        }
        found
    }

    fn optimal_k_num(bitmap_bits: u64, items_count: uint) -> uint {
        let m = bitmap_bits as f64;
        let n = items_count as f64;
        let k_num = (m / n * num::ln(2.0f64).ceil()) as uint;
        cmp::max(k_num, 1)
    }

    fn bloom_hash<T: Hash>(&self, hashes: & mut [u64, ..2],
                  item: &T, k_i: uint) -> u64 {
        if k_i < 2 {
            let skey = self.skeys[k_i];
            let hash = (*item).hash_keyed(skey.k1, skey.k2);
            hashes[k_i] = hash;
            hash
        } else {
            hashes[0] + (((k_i as u64) * hashes[1]) % 0xffffffffffffffc5)
        }
    }
}

#[test]
fn bloom_test_set() {
    let mut bloom = Bloom::new(10, 80);
    let key = &rand::task_rng().gen_ascii_str(16u);
    assert!(bloom.check(key) == false);
    bloom.set(&key);
    assert!(bloom.check(key.clone()) == true);
}

#[test]
fn bloom_test_check_and_set() {
    let mut bloom = Bloom::new(10, 80);
    let key = &rand::task_rng().gen_ascii_str(16u);
    assert!(bloom.check_and_set(key) == false);
    assert!(bloom.check_and_set(key.clone()) == true);
}
<commit_msg>Move to std::hash::sip<commit_after>// (C)opyleft 2013 Frank Denis

/*!
 * Bloom filter for Rust
 *
 * This is a simple but fast Bloom filter implementation, that requires only
 * 2 hash functions, generated with SipHash-2-4 using randomized keys.
 */

#[desc = "A fast Bloom filter implementation."];
#[license = "BSD"];
#[crate_id="bloomfilter#0.1"];
#[crate_type = "rlib"];
#[warn(non_camel_case_types, non_uppercase_statics, non_uppercase_statics, unnecessary_qualification, managed_heap_memory)];

extern crate collections;
extern crate extra;

use std::cmp;
use std::hash::Hash;
use std::num;
use std::rand;
use std::hash::sip;
use std::rand::Rng;
use collections::bitv;

struct SipHashKey {
    k1: u64,
    k2: u64
}

impl SipHashKey {
    fn new_random() -> SipHashKey {
        let mut rng = rand::task_rng();
        SipHashKey {
            k1: rand::Rand::rand(& mut rng),
            k2: rand::Rand::rand(& mut rng)
        }
    }
}

/// Bloom filter structure
pub struct Bloom {
    priv bitmap: bitv::Bitv,
    priv bitmap_bits: u64,
    priv k_num: uint,
    priv skeys: [SipHashKey, ..2]
}

impl Bloom {

/// Create a new bloom filter structure.
/// bitmap_size is the size in bytes (not bits) that will be allocated in memory
/// items_count is an estimation of the maximum number of items to store.
    pub fn new(bitmap_size: uint, items_count: uint) -> Bloom {
        assert!(bitmap_size > 0u && items_count > 0u);
        let bitmap_bits = (bitmap_size as u64) * 8u64;
        let k_num = Bloom::optimal_k_num(bitmap_bits, items_count);
        let bitmap = bitv::Bitv::new(bitmap_bits as uint, false);
        let skeys = [ SipHashKey::new_random(), SipHashKey::new_random() ];
        Bloom {
            bitmap: bitmap,
            bitmap_bits: bitmap_bits,
            k_num: k_num,
            skeys: skeys
        }
    }

/// Create a new bloom filter structure.
/// items_count is an estimation of the maximum number of items to store.
/// fp_p is the wanted rate of false positives, in ]0.0, 1.0[
    pub fn new_for_fp_rate(items_count: uint, fp_p: f64) -> Bloom {
        let bitmap_size = Bloom::compute_bitmap_size(items_count, fp_p);
        Bloom::new(bitmap_size, items_count)
    }

/// Compute a recommended bitmap size for items_count items
/// and a fp_p rate of false positives.
/// fp_p obviously has to be within the ]0.0, 1.0[ range.
    pub fn compute_bitmap_size(items_count: uint, fp_p: f64) -> uint {
        assert!(items_count > 0u);
        assert!(fp_p > 0.0 && fp_p < 1.0);
        let log2 = num::ln(2.0f64);
        let log2_2 = log2 * log2;
        ((items_count as f64) * num::ln(fp_p) / (-8.0 * log2_2)).ceil() as uint
    }

/// Record the presence of an item.
    pub fn set<T: Hash>(& mut self, item: T) {
        let mut hashes = [ 0u64, 0u64 ];
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            self.bitmap.set(bit_offset, true);
        }
    }

/// Check if an item is present in the set.
/// There can be false positives, but no false negatives.
    pub fn check<T: Hash>(&self, item: T) -> bool {
        let mut hashes = [ 0u64, 0u64 ];
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            if self.bitmap.get(bit_offset) == false {
                return false;
            }
        }
        true
    }

/// Record the presence of an item in the set,
/// and return the previous state of this item.
    pub fn check_and_set<T: Hash>(&mut self, item: T) -> bool {
        let mut hashes = [ 0u64, 0u64 ];
        let mut found = true;
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            if self.bitmap.get(bit_offset) == false {
                found = false;
                self.bitmap.set(bit_offset, true);
            }
        }
        found
    }

    fn optimal_k_num(bitmap_bits: u64, items_count: uint) -> uint {
        let m = bitmap_bits as f64;
        let n = items_count as f64;
        let k_num = (m / n * num::ln(2.0f64).ceil()) as uint;
        cmp::max(k_num, 1)
    }

    fn bloom_hash<T: Hash>(&self, hashes: & mut [u64, ..2],
                  item: &T, k_i: uint) -> u64 {
        if k_i < 2 {
            let skey = self.skeys[k_i];
            let hash = sip::hash_with_keys(skey.k1, skey.k2, item);
            hashes[k_i] = hash;
            hash
        } else {
            hashes[0] + (((k_i as u64) * hashes[1]) % 0xffffffffffffffc5)
        }
    }
}

#[test]
fn bloom_test_set() {
    let mut bloom = Bloom::new(10, 80);
    let key = &rand::task_rng().gen_ascii_str(16u);
    assert!(bloom.check(key) == false);
    bloom.set(&key);
    assert!(bloom.check(key.clone()) == true);
}

#[test]
fn bloom_test_check_and_set() {
    let mut bloom = Bloom::new(10, 80);
    let key = &rand::task_rng().gen_ascii_str(16u);
    assert!(bloom.check_and_set(key) == false);
    assert!(bloom.check_and_set(key.clone()) == true);
}
<|endoftext|>