<commit_before>//! emit_bosun -- Emits a single metric datum and corresponding metric meta data to [Bosun](http:////bosun.org).

extern crate clap;
extern crate env_logger;
#[macro_use]
extern crate log;

extern crate bosun_emitter;

use clap::{Arg, ArgMatches, App};
use std::error::Error;
use std::path::Path;

use bosun_emitter::{BosunClient, Metadata, Datum, Tags, EmitterError, BosunConfig};

static VERSION: &'static str = env!("CARGO_PKG_VERSION");
static DEFAULT_CONFIG_FILE: &'static str = "/etc/bosun/scollector.conf";

#[derive(Debug)]
struct Config {
    host: String,
    hostname: String,
    metric: Option<String>,
    value: Option<String>,
    rate: Option<String>,
    unit: Option<String>,
    description: Option<String>,
    tags: Tags,
}

impl Config {
    pub fn default() -> Config {
        let bosun_config = BosunConfig::default();
        Config::from_bosun_config(bosun_config)
    }

    pub fn from_bosun_config(bosun_config: BosunConfig) -> Config {
        Config {
            host: bosun_config.Host,
            hostname: bosun_config.Hostname,
            metric: None,
            value: None,
            rate: None,
            unit: None,
            description: None,
            tags: bosun_config.Tags
        }
    }
}

fn main() {
    if env_logger::init().is_err() {
        exit_with_error("Could not initiliaze logger", -1);
    }

    let app = App::new("Emit a Bosun")
                       .version(VERSION)
                       .after_help("Two modes are supported, i.e., sending a datum with meta \
                                    data or sending only meta data. The modes are controlled \
                                    whether a value `--value` is passed or not. Please mind that \
                                    in both cases the meta data is required.")
                       .arg(Arg::with_name("config")
                                .short("c")
                                .long("config")
                                .value_name("FILE")
                                .help("Sets a custom config file")
                                .takes_value(true))
                       .arg(Arg::with_name("host")
                                .long("host")
                                .value_name("HOST:PORT")
                                .help("Sets Bosun server connection parameters")
                                .takes_value(true))
                       .arg(Arg::with_name("hostname")
                                .long("hostname")
                                .value_name("HOSTNAME")
                                .help("Sets hostname")
                                .takes_value(true))
                       .arg(Arg::with_name("metric")
                                .short("m")
                                .long("metric")
                                .value_name("METRIC NAME")
                                .help("Sets metric name")
                                .takes_value(true))
                       .arg(Arg::with_name("value")
                                .requires("metric")
                                .short("v")
                                .long("value")
                                .value_name("VALUE")
                                .help("Sets metric value")
                                .takes_value(true))
                       .arg(Arg::with_name("rate")
                                .requires_all(&["metric", "rate", "unit", "description"])
                                .short("r")
                                .long("rate")
                                .value_name("RATE")
                                .possible_values(&["gauge", "counter", "rate"])
                                .help("Sets rate type")
                                .takes_value(true))
                       .arg(Arg::with_name("unit")
                                .requires_all(&["metric", "rate", "unit", "description"])
                                .short("u")
                                .long("unit")
                                .value_name("UNIT")
                                .help("Sets metric value unit")
                                .takes_value(true))
                       .arg(Arg::with_name("description")
                                .requires_all(&["metric", "rate", "unit", "description"])
                                .short("d")
                                .long("description")
                                .value_name("DESCRIPTION")
                                .help("Sets metric description")
                                .takes_value(true))
                       .arg(Arg::with_name("tags")
                                .use_delimiter(false)
                                .short("t")
                                .long("tags")
                                .value_name("KEY1=VALUE1,KEY2=VALUE2,...")
                                .help("Sets tags")
                                .takes_value(true))
                       .arg(Arg::with_name("show-config")
                                .long("show-config")
                                .help("Prints config"))
                       .arg(Arg::with_name("verbose")
                                .long("verbose")
                                .help("Enables verbose output"))
                       .arg(Arg::with_name("force")
                                .hidden(true)
                                .long("force")
                                .help("Forces metric datum to be send even without meta data"));
    let cli_args = app.get_matches();

    let force: bool = cli_args.is_present("force");
    let verbose: bool = cli_args.is_present("verbose");
    let config: Config = match parse_args(&cli_args) {
        Ok(config) => config,
        Err(err) => {
            exit_with_error(&format!("Failed to parse configuration, because {}.", err),
                            -2);
        }
    };
    if cli_args.is_present("show-config") {
        println!("config: {:?}", config);
    }

    let mode = match mode(&config, force) {
        Ok(mode) => mode,
        Err(ModeError::NoMetadata) => exit_with_error("Cannot send datum without meta data.", -11),
        Err(ModeError::NoValue) => exit_with_error("Cannot send datum without value.", -12),
        Err(ModeError::NoSuchMode) => {
            println!("Command line arguments combination does not make any sense.\n\n{}", cli_args.usage());
            exit_with_error("", -13);
        }
    };

    let result = run(&config, mode, verbose);
    match result {
        Ok(_) => {}
        Err(EmitterError::JsonParseError(e)) => {
            exit_with_error(&format!("Failed to create JSON document, because {}.", e),
                            1)
        }
        Err(EmitterError::EmitError(e)) => {
            exit_with_error(&format!("Failed to send, because {}.", e), 2)
        }
        Err(EmitterError::ReceiveError(e)) => {
            exit_with_error(&format!("Failed to create resource, because {}.", e), 3)
        }
    }
}

fn parse_args(cli_args: &ArgMatches) -> Result<Config, Box<Error>> {
    let bosun_config_file_path = Path::new(cli_args.value_of("config").unwrap_or(DEFAULT_CONFIG_FILE));
    let mut config: Config = if bosun_config_file_path.exists() {
        let bosun_config: BosunConfig = try!(BosunConfig::load_from_scollector_config(&bosun_config_file_path));
        Config::from_bosun_config(bosun_config)
    } else {
        Config::default()
    };

    if cli_args.is_present("host") {
        config.host = cli_args.value_of("host").unwrap().to_string();
    }

    config.tags.insert("host".to_string(), config.hostname.to_string());
    if cli_args.is_present("hostname") {
        config.hostname = cli_args.value_of("hostname").unwrap().to_string();
        config.tags.insert("host".to_string(), config.hostname.to_string());
    }

    if cli_args.is_present("metric") {
        config.metric = Some(cli_args.value_of("metric").unwrap().to_string());
    }

    if cli_args.is_present("value") {
        config.value = Some(cli_args.value_of("value").unwrap().to_string());
    }

    if cli_args.is_present("rate") {
        config.rate = Some(cli_args.value_of("rate").unwrap().to_string());
    }

    if cli_args.is_present("unit") {
        config.unit = Some(cli_args.value_of("unit").unwrap().to_string());
    }

    if cli_args.is_present("description") {
        config.description = Some(cli_args.value_of("description").unwrap().to_string());
    }

    if cli_args.is_present("tags") {
        let tags_string = cli_args.value_of("tags").unwrap().to_string();
        try!(parse_tags(&mut config, &tags_string));
    }

    Ok(config)
}

fn parse_tags(config: &mut Config, tags_string: &str) -> Result<(), String> {
    let tags = tags_string.split(',');
    for tag in tags {
        let kv = tag.split('=').collect::<Vec<&str>>();
        if kv.len() != 2 {
            return Err(format!("unable to parse tags: '{}'", tags_string));
        }
        let k = kv[0].to_string();
        let v = kv[1].to_string();
        config.tags.insert(k, v);
    }

    Ok(())
}

enum Mode {
    Normal,
    MetadataOnly,
    DatumOnly,
}

enum ModeError {
    NoMetadata,
    NoValue,
    NoSuchMode,
}

fn run(config: &Config, mode: Mode, verbose: bool) -> bosun_emitter::EmitterResult {
    match mode {
        Mode::Normal => {
            msg("Sending meta data.", verbose);
            try!(emit_metadata(config));

            msg("Sending datum.", verbose);
            emit_datum(config)
        }
        Mode::MetadataOnly => {
           msg("Sending meta data.", verbose);
            emit_metadata(config)
        }
        Mode::DatumOnly => {
            msg("Sending datum.", verbose);
            emit_datum(config)
        }
    }
}

/// We support two modes officially and more mode unofficially.
/// 1. Send Datum with Metadata
/// 1. Send only Metadata
/// 1. Send Datum without Metadata -- only with `--force`
fn mode(config: &Config, force: bool) -> Result<Mode, ModeError> {
    let mode_config = (config.metric.is_some(),
                       config.value.is_some(),
                       config.rate.is_some(),
                       config.unit.is_some(),
                       config.description.is_some(),
                       force);
    match mode_config {
        (true, true, true, true, true, _) => Ok(Mode::Normal),
        (true, _, true, true, true, _) => Ok(Mode::MetadataOnly),
        (true, true, _, _, _, true) => Ok(Mode::DatumOnly),
        (true, true, _, _, _, false) => Err(ModeError::NoMetadata),
        (true, false, _, _, _, _) => Err(ModeError::NoValue),
        _ => Err(ModeError::NoSuchMode),
    }
}

fn msg(msg: &str, verbose: bool) {
    if verbose {
        println!("{}", msg);
    }
}

fn exit_with_error(msg: &str, exit_code: i32) -> ! {
    println!("{}", msg);
    std::process::exit(exit_code);
}

fn emit_datum(config: &Config) -> bosun_emitter::EmitterResult {
    // unwraps are safe, because mode analysis already checked these values are set
    let client = BosunClient::new(&config.host);
    let datum = Datum::now(config.metric.as_ref().unwrap(),
                           config.value.as_ref().unwrap(),
                           &config.tags);
    client.emit_datum(&datum)
}

fn emit_metadata(config: &Config) -> bosun_emitter::EmitterResult {
    // unwraps are safe, because mode analysis already checked these values are set
    let client = BosunClient::new(&config.host);
    let metadata = Metadata::new(config.metric.as_ref().unwrap(),
                                 config.rate.as_ref().unwrap(),
                                 config.unit.as_ref().unwrap(),
                                 config.description.as_ref().unwrap());
    client.emit_metadata(&metadata)
}

#[cfg(test)]
mod tests {
    use super::{Config, parse_tags};

    #[test]
    fn parse_tags_test_okay() {
        let mut config = Config::default();
        let tags = "key1=val1,key2=val2";
        let _ = parse_tags(&mut config, &tags);
        assert_eq!(config.tags.len(), 2);
    }

    #[test]
    #[should_panic(expected = "assertion failed")]
    fn parse_tags_test_fails_wrong_kv_separator() {
        let mut config = Config::default();
        let tags = "key1=val1,key2:val2";
        let _ = parse_tags(&mut config, &tags);
        assert_eq!(config.tags.len(), 2);
    }
}

<commit_msg>Fix minor compiler warning<commit_after>//! emit_bosun -- Emits a single metric datum and corresponding metric meta data to [Bosun](http:////bosun.org).

extern crate clap;
extern crate env_logger;
extern crate log;

extern crate bosun_emitter;

use clap::{Arg, ArgMatches, App};
use std::error::Error;
use std::path::Path;

use bosun_emitter::{BosunClient, Metadata, Datum, Tags, EmitterError, BosunConfig};

static VERSION: &'static str = env!("CARGO_PKG_VERSION");
static DEFAULT_CONFIG_FILE: &'static str = "/etc/bosun/scollector.conf";

#[derive(Debug)]
struct Config {
    host: String,
    hostname: String,
    metric: Option<String>,
    value: Option<String>,
    rate: Option<String>,
    unit: Option<String>,
    description: Option<String>,
    tags: Tags,
}

impl Config {
    pub fn default() -> Config {
        let bosun_config = BosunConfig::default();
        Config::from_bosun_config(bosun_config)
    }

    pub fn from_bosun_config(bosun_config: BosunConfig) -> Config {
        Config {
            host: bosun_config.Host,
            hostname: bosun_config.Hostname,
            metric: None,
            value: None,
            rate: None,
            unit: None,
            description: None,
            tags: bosun_config.Tags
        }
    }
}

fn main() {
    if env_logger::init().is_err() {
        exit_with_error("Could not initiliaze logger", -1);
    }

    let app = App::new("Emit a Bosun")
                       .version(VERSION)
                       .after_help("Two modes are supported, i.e., sending a datum with meta \
                                    data or sending only meta data. The modes are controlled \
                                    whether a value `--value` is passed or not. Please mind that \
                                    in both cases the meta data is required.")
                       .arg(Arg::with_name("config")
                                .short("c")
                                .long("config")
                                .value_name("FILE")
                                .help("Sets a custom config file")
                                .takes_value(true))
                       .arg(Arg::with_name("host")
                                .long("host")
                                .value_name("HOST:PORT")
                                .help("Sets Bosun server connection parameters")
                                .takes_value(true))
                       .arg(Arg::with_name("hostname")
                                .long("hostname")
                                .value_name("HOSTNAME")
                                .help("Sets hostname")
                                .takes_value(true))
                       .arg(Arg::with_name("metric")
                                .short("m")
                                .long("metric")
                                .value_name("METRIC NAME")
                                .help("Sets metric name")
                                .takes_value(true))
                       .arg(Arg::with_name("value")
                                .requires("metric")
                                .short("v")
                                .long("value")
                                .value_name("VALUE")
                                .help("Sets metric value")
                                .takes_value(true))
                       .arg(Arg::with_name("rate")
                                .requires_all(&["metric", "rate", "unit", "description"])
                                .short("r")
                                .long("rate")
                                .value_name("RATE")
                                .possible_values(&["gauge", "counter", "rate"])
                                .help("Sets rate type")
                                .takes_value(true))
                       .arg(Arg::with_name("unit")
                                .requires_all(&["metric", "rate", "unit", "description"])
                                .short("u")
                                .long("unit")
                                .value_name("UNIT")
                                .help("Sets metric value unit")
                                .takes_value(true))
                       .arg(Arg::with_name("description")
                                .requires_all(&["metric", "rate", "unit", "description"])
                                .short("d")
                                .long("description")
                                .value_name("DESCRIPTION")
                                .help("Sets metric description")
                                .takes_value(true))
                       .arg(Arg::with_name("tags")
                                .use_delimiter(false)
                                .short("t")
                                .long("tags")
                                .value_name("KEY1=VALUE1,KEY2=VALUE2,...")
                                .help("Sets tags")
                                .takes_value(true))
                       .arg(Arg::with_name("show-config")
                                .long("show-config")
                                .help("Prints config"))
                       .arg(Arg::with_name("verbose")
                                .long("verbose")
                                .help("Enables verbose output"))
                       .arg(Arg::with_name("force")
                                .hidden(true)
                                .long("force")
                                .help("Forces metric datum to be send even without meta data"));
    let cli_args = app.get_matches();

    let force: bool = cli_args.is_present("force");
    let verbose: bool = cli_args.is_present("verbose");
    let config: Config = match parse_args(&cli_args) {
        Ok(config) => config,
        Err(err) => {
            exit_with_error(&format!("Failed to parse configuration, because {}.", err),
                            -2);
        }
    };
    if cli_args.is_present("show-config") {
        println!("config: {:?}", config);
    }

    let mode = match mode(&config, force) {
        Ok(mode) => mode,
        Err(ModeError::NoMetadata) => exit_with_error("Cannot send datum without meta data.", -11),
        Err(ModeError::NoValue) => exit_with_error("Cannot send datum without value.", -12),
        Err(ModeError::NoSuchMode) => {
            println!("Command line arguments combination does not make any sense.\n\n{}", cli_args.usage());
            exit_with_error("", -13);
        }
    };

    let result = run(&config, mode, verbose);
    match result {
        Ok(_) => {}
        Err(EmitterError::JsonParseError(e)) => {
            exit_with_error(&format!("Failed to create JSON document, because {}.", e),
                            1)
        }
        Err(EmitterError::EmitError(e)) => {
            exit_with_error(&format!("Failed to send, because {}.", e), 2)
        }
        Err(EmitterError::ReceiveError(e)) => {
            exit_with_error(&format!("Failed to create resource, because {}.", e), 3)
        }
    }
}

fn parse_args(cli_args: &ArgMatches) -> Result<Config, Box<Error>> {
    let bosun_config_file_path = Path::new(cli_args.value_of("config").unwrap_or(DEFAULT_CONFIG_FILE));
    let mut config: Config = if bosun_config_file_path.exists() {
        let bosun_config: BosunConfig = try!(BosunConfig::load_from_scollector_config(&bosun_config_file_path));
        Config::from_bosun_config(bosun_config)
    } else {
        Config::default()
    };

    if cli_args.is_present("host") {
        config.host = cli_args.value_of("host").unwrap().to_string();
    }

    config.tags.insert("host".to_string(), config.hostname.to_string());
    if cli_args.is_present("hostname") {
        config.hostname = cli_args.value_of("hostname").unwrap().to_string();
        config.tags.insert("host".to_string(), config.hostname.to_string());
    }

    if cli_args.is_present("metric") {
        config.metric = Some(cli_args.value_of("metric").unwrap().to_string());
    }

    if cli_args.is_present("value") {
        config.value = Some(cli_args.value_of("value").unwrap().to_string());
    }

    if cli_args.is_present("rate") {
        config.rate = Some(cli_args.value_of("rate").unwrap().to_string());
    }

    if cli_args.is_present("unit") {
        config.unit = Some(cli_args.value_of("unit").unwrap().to_string());
    }

    if cli_args.is_present("description") {
        config.description = Some(cli_args.value_of("description").unwrap().to_string());
    }

    if cli_args.is_present("tags") {
        let tags_string = cli_args.value_of("tags").unwrap().to_string();
        try!(parse_tags(&mut config, &tags_string));
    }

    Ok(config)
}

fn parse_tags(config: &mut Config, tags_string: &str) -> Result<(), String> {
    let tags = tags_string.split(',');
    for tag in tags {
        let kv = tag.split('=').collect::<Vec<&str>>();
        if kv.len() != 2 {
            return Err(format!("unable to parse tags: '{}'", tags_string));
        }
        let k = kv[0].to_string();
        let v = kv[1].to_string();
        config.tags.insert(k, v);
    }

    Ok(())
}

enum Mode {
    Normal,
    MetadataOnly,
    DatumOnly,
}

enum ModeError {
    NoMetadata,
    NoValue,
    NoSuchMode,
}

fn run(config: &Config, mode: Mode, verbose: bool) -> bosun_emitter::EmitterResult {
    match mode {
        Mode::Normal => {
            msg("Sending meta data.", verbose);
            try!(emit_metadata(config));

            msg("Sending datum.", verbose);
            emit_datum(config)
        }
        Mode::MetadataOnly => {
           msg("Sending meta data.", verbose);
            emit_metadata(config)
        }
        Mode::DatumOnly => {
            msg("Sending datum.", verbose);
            emit_datum(config)
        }
    }
}

/// We support two modes officially and more mode unofficially.
/// 1. Send Datum with Metadata
/// 1. Send only Metadata
/// 1. Send Datum without Metadata -- only with `--force`
fn mode(config: &Config, force: bool) -> Result<Mode, ModeError> {
    let mode_config = (config.metric.is_some(),
                       config.value.is_some(),
                       config.rate.is_some(),
                       config.unit.is_some(),
                       config.description.is_some(),
                       force);
    match mode_config {
        (true, true, true, true, true, _) => Ok(Mode::Normal),
        (true, _, true, true, true, _) => Ok(Mode::MetadataOnly),
        (true, true, _, _, _, true) => Ok(Mode::DatumOnly),
        (true, true, _, _, _, false) => Err(ModeError::NoMetadata),
        (true, false, _, _, _, _) => Err(ModeError::NoValue),
        _ => Err(ModeError::NoSuchMode),
    }
}

fn msg(msg: &str, verbose: bool) {
    if verbose {
        println!("{}", msg);
    }
}

fn exit_with_error(msg: &str, exit_code: i32) -> ! {
    println!("{}", msg);
    std::process::exit(exit_code);
}

fn emit_datum(config: &Config) -> bosun_emitter::EmitterResult {
    // unwraps are safe, because mode analysis already checked these values are set
    let client = BosunClient::new(&config.host);
    let datum = Datum::now(config.metric.as_ref().unwrap(),
                           config.value.as_ref().unwrap(),
                           &config.tags);
    client.emit_datum(&datum)
}

fn emit_metadata(config: &Config) -> bosun_emitter::EmitterResult {
    // unwraps are safe, because mode analysis already checked these values are set
    let client = BosunClient::new(&config.host);
    let metadata = Metadata::new(config.metric.as_ref().unwrap(),
                                 config.rate.as_ref().unwrap(),
                                 config.unit.as_ref().unwrap(),
                                 config.description.as_ref().unwrap());
    client.emit_metadata(&metadata)
}

#[cfg(test)]
mod tests {
    use super::{Config, parse_tags};

    #[test]
    fn parse_tags_test_okay() {
        let mut config = Config::default();
        let tags = "key1=val1,key2=val2";
        let _ = parse_tags(&mut config, &tags);
        assert_eq!(config.tags.len(), 2);
    }

    #[test]
    #[should_panic(expected = "assertion failed")]
    fn parse_tags_test_fails_wrong_kv_separator() {
        let mut config = Config::default();
        let tags = "key1=val1,key2:val2";
        let _ = parse_tags(&mut config, &tags);
        assert_eq!(config.tags.len(), 2);
    }
}

<|endoftext|>


<commit_before>//! frust
//! A find replacement with SQL-like syntax.
//!
//! # TODOs
//!    - [ ] if query does not end with ';' or is empty, append a ';'
//!    - [X] implement the --depth option
//!

extern crate clap;
extern crate walkdir;
extern crate frustlib;
extern crate regex;
// #[macro_use] extern crate log;
// extern crate env_logger;

use regex::Regex;
use clap::{App, Arg};
use frustlib::query::Query;

fn is_integer(inp: String) -> Result<(), String> {
    match inp.parse::<u32>() {
        Ok(_) => Ok(()),
        _     => Err(String::from("The argument does not seem to be an unsigned integer.")),
    }
}

fn main() {
    let matches = App::new("frust")
        .version("0.0.1")
        .about("Does great stuff, in the future, hopefully...")
        .author("Timm Behner, Martin Clauß")
        .arg(Arg::with_name("QUERY")
             .help("find files according to the query the directory tree")
             .index(1)
             .required(false)
         )
        .arg(Arg::with_name("depth")
             .short("d")
             .long("depth")
             .help("maximum depth of search")
             .required(false)
             .takes_value(true)
             .value_name("DEPTH")
             .default_value("4096") // current maximum directory tree depth on my linux machine
             .validator(is_integer)
         )
        .arg(Arg::with_name("machine-readable")
             .short("m")
             .long("machine-readable")
             .help("print attributes in machine readable syntax")
             .required(false)
             .takes_value(false)
        )
        .get_matches();

    let mut q = if matches.is_present("QUERY") {
        let mut query_string = String::from(matches.value_of("QUERY").unwrap());
        let semicolon_check = Regex::new(";\\s*$").unwrap();
        if !semicolon_check.is_match(&query_string) {
            query_string.push(';');
        }
        Query::parse(&query_string)
    } else {
        panic!("frust without a query is not yet supported!");
    };

    let max_depth = matches.value_of("depth").unwrap().parse::<usize>().unwrap();
    let machine_mode = matches.is_present("machine-readable");
    q.execute(max_depth, machine_mode);
}
<commit_msg>Refactor<commit_after>//! frust
//! A find replacement with SQL-like syntax.
//!
//! # TODOs
//!    - [ ] if query does not end with ';' or is empty, append a ';'
//!    - [X] implement the --depth option
//!

extern crate clap;
extern crate walkdir;
extern crate frustlib;
extern crate regex;
// #[macro_use] extern crate log;
// extern crate env_logger;

use regex::Regex;
use clap::{App, Arg};
use frustlib::query::Query;

fn is_integer(inp: String) -> Result<(), String> {
    match inp.parse::<u32>() {
        Ok(_) => Ok(()),
        _     => Err(String::from("The argument does not seem to be an unsigned integer.")),
    }
}

fn check_semicolon(inp: &mut String){
    let semicolon_check = Regex::new(";\\s*$").unwrap();
    if !semicolon_check.is_match(inp) {
        inp.push(';');
    }
}

fn main() {
    let matches = App::new("frust")
        .version("0.0.1")
        .about("Does great stuff, in the future, hopefully...")
        .author("Timm Behner, Martin Clauß")
        .arg(Arg::with_name("QUERY")
             .help("find files according to the query the directory tree")
             .index(1)
             .required(false)
         )
        .arg(Arg::with_name("depth")
             .short("d")
             .long("depth")
             .help("maximum depth of search")
             .required(false)
             .takes_value(true)
             .value_name("DEPTH")
             .default_value("4096") // current maximum directory tree depth on my linux machine
             .validator(is_integer)
         )
        .arg(Arg::with_name("machine-readable")
             .short("m")
             .long("machine-readable")
             .help("print attributes in machine readable syntax")
             .required(false)
             .takes_value(false)
        )
        .get_matches();

    let mut q = match matches.value_of("QUERY") {
        Some(query_inp) => {
            let mut query_string = String::from(query_inp);
            check_semicolon(&mut query_string);
            Query::parse(&query_string)
        },
        None => panic!("frust without a query is not yet supported!"),
    };

    let max_depth = matches.value_of("depth").unwrap().parse::<usize>().expect("Given depth cannot be parsed to an integer!");
    let machine_mode = matches.is_present("machine-readable");
    q.execute(max_depth, machine_mode);
}
<|endoftext|>


<commit_before>extern crate iron;
extern crate router;
extern crate rustc_serialize;

use std::io::Read;
use std::sync::{Arc, Mutex};
use std::collections::HashMap;

use iron::prelude::*;
use iron::status;
use router::Router;
use rustc_serialize::json::{self, Json};

#[derive(Debug)]
struct Document {
    data: Json,
}

impl Document {
    fn from_json(data: Json) -> Document {
        Document{
            data: data,
        }
    }
}

#[derive(Debug)]
enum Filter {
    Not(Box<Filter>),
    Or(Vec<Filter>),
    And(Vec<Filter>),
    Term(String, String),
}

impl Filter {
    fn matches(&self, doc: &Document) -> bool {
        match *self {
            Filter::Not(ref filter) => !filter.matches(doc),
            Filter::Or(ref filters) => {
                for filter in filters.iter() {
                    if (filter.matches(doc)) {
                        return true;
                    }
                }

                false
            },
            Filter::And(ref filters) => {
                for filter in filters.iter() {
                    if (!filter.matches(doc)) {
                        return false;
                    }
                }

                true
            },
            Filter::Term(ref field, ref value) => {
                let obj = doc.data.as_object().unwrap();

                if let Some(field_value) = obj.get(field) {
                    if let Json::String(ref field_value) = *field_value {
                        return field_value == value
                    }
                }

                false
            }
        }
    }
}

#[derive(Debug)]
struct Mapping {
    pub docs: HashMap<String, Document>,
}

impl Mapping {
    fn new() -> Mapping {
        Mapping{
            docs: HashMap::new(),
        }
    }
}

#[derive(Debug)]
struct Index {
    pub mappings: HashMap<String, Mapping>,
}

impl Index {
    fn new() -> Index {
        Index{
            mappings: HashMap::new(),
        }
    }
}

fn index_not_found_response() -> Response {
    let mut response = Response::with((status::NotFound, "{\"message\": \"Index not found\"}"));
    response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
    return response;
}

fn main() {
    let indices = Arc::new(Mutex::new(HashMap::new()));
    let mut wagtail_index = Index::new();
    wagtail_index.mappings.insert("wagtaildocs_document".to_owned(), Mapping::new());
    indices.lock().unwrap().insert("wagtaildemo".to_owned(), wagtail_index);

    let f = Filter::Or(vec![
        Filter::Term("title".to_owned(), "test".to_owned()),
        Filter::Term("title".to_owned(), "foo".to_owned()),
    ]);

    println!("{:?}", f);

    let mut router = Router::new();

    router.get("/", |_: &mut Request| -> IronResult<Response> {
        Ok(Response::with((status::Ok, "Hello World!")))
    });

    {
        let indices = indices.clone();

        router.get("/:index/_count", move |req: &mut Request| -> IronResult<Response> {
            // URL parameters
            let index_name = req.extensions.get::<Router>().unwrap().find("index").unwrap_or("");

            // Lock index array
            let indices = indices.lock().unwrap();

            // Find index
            let index = match indices.get(index_name) {
                Some(index) => index,
                None => {
                    return Ok(index_not_found_response());
                }
            };

            // Load query from body
            let mut payload = String::new();
            req.body.read_to_string(&mut payload).unwrap();

            let data = if !payload.is_empty() {
                Some(match Json::from_str(&payload) {
                    Ok(data) => data,
                    Err(error) => {
                        // TODO: What specifically is bad about the JSON?
                        let mut response = Response::with((status::BadRequest, "{\"message\": \"Couldn't parse JSON\"}"));
                        response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
                        return Ok(response);
                    }
                })
            } else {
                None
            };

            // TODO: Run query

            // Temporary count and return numbers
            let mut count = 0;
            for mapping in index.mappings.values() {
                count += mapping.docs.len();
            }

            let mut response = Response::with((status::Ok, format!("{{\"count\": {}}}", count)));
            response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
            Ok(response)
        });
    }

    {
        let indices = indices.clone();

        router.get("/:index/_search", move |req: &mut Request| -> IronResult<Response> {
            // URL parameters
            let index_name = req.extensions.get::<Router>().unwrap().find("index").unwrap_or("");

            // Lock index array
            let indices = indices.lock().unwrap();

            // Find index
            let index = match indices.get(index_name) {
                Some(index) => index,
                None => {
                    return Ok(index_not_found_response());
                }
            };

            // Load query from body
            let mut payload = String::new();
            req.body.read_to_string(&mut payload).unwrap();

            let data = if !payload.is_empty() {
                Some(match Json::from_str(&payload) {
                    Ok(data) => data,
                    Err(error) => {
                        // TODO: What specifically is bad about the JSON?
                        let mut response = Response::with((status::BadRequest, "{\"message\": \"Couldn't parse JSON\"}"));
                        response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
                        return Ok(response);
                    }
                })
            } else {
                None
            };

            // TODO: Run query

            let mut response = Response::with((status::Ok, "{\"hits\": {\"total\": 0, \"hits\": []}}"));
            response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
            Ok(response)
        });
    }

    {
        let indices = indices.clone();

        router.get("/:index/:mapping/:doc", move |req: &mut Request| -> IronResult<Response> {
            // URL parameters
            let index_name = req.extensions.get::<Router>().unwrap().find("index").unwrap_or("");
            let mapping_name = req.extensions.get::<Router>().unwrap().find("mapping").unwrap_or("");
            let doc_id = req.extensions.get::<Router>().unwrap().find("doc").unwrap_or("");

            // Lock index array
            let indices = indices.lock().unwrap();

            // Find index
            let index = match indices.get(index_name) {
                Some(index) => index,
                None => {
                    return Ok(index_not_found_response());
                }
            };

            // Find mapping
            let mapping = match index.mappings.get(mapping_name) {
                Some(mapping) => mapping,
                None => {
                    let mut response = Response::with((status::NotFound, "{\"message\": \"Mapping not found\"}"));
                    response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
                    return Ok(response);
                }
            };

            // Find document
            let doc = match mapping.docs.get(doc_id) {
                Some(doc) => doc,
                None => {
                    let mut response = Response::with((status::NotFound, "{\"message\": \"Document not found\"}"));
                    response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
                    return Ok(response);
                }
            };

            let mut response = Response::with((status::Ok, json::encode(&doc.data).unwrap()));
            response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
            Ok(response)
        });
    }

    {
        let indices = indices.clone();

        router.put("/:index/:mapping/:doc", move |req: &mut Request| -> IronResult<Response> {
            // URL parameters
            let index_name = req.extensions.get::<Router>().unwrap().find("index").unwrap_or("");
            let mapping_name = req.extensions.get::<Router>().unwrap().find("mapping").unwrap_or("");
            let ref doc_id = req.extensions.get::<Router>().unwrap().find("doc").unwrap_or("");

            // Lock index array
            let mut indices = indices.lock().unwrap();

            // Find index
            let mut index = match indices.get_mut(index_name) {
                Some(index) => index,
                None => {
                    return Ok(index_not_found_response());
                }
            };

            // Find mapping
            let mut mapping = match index.mappings.get_mut(mapping_name) {
                Some(mapping) => mapping,
                None => {
                    let mut response = Response::with((status::NotFound, "{\"message\": \"Mapping not found\"}"));
                    response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
                    return Ok(response);
                }
            };

            // Load data from body
            let mut payload = String::new();
            req.body.read_to_string(&mut payload).unwrap();

            let data = if !payload.is_empty() {
                Some(match Json::from_str(&payload) {
                    Ok(data) => data,
                    Err(error) => {
                        // TODO: What specifically is bad about the JSON?
                        let mut response = Response::with((status::BadRequest, "{\"message\": \"Couldn't parse JSON\"}"));
                        response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
                        return Ok(response);
                    }
                })
            } else {
                None
            };

            // Create and insert document
            if let Some(data) = data {
                let doc = Document::from_json(data);
                println!("{:?}", f.matches(&doc));
                mapping.docs.insert(doc_id.clone().to_owned(), doc);
            }

            let mut response = Response::with((status::Ok, "{}"));
            response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
            Ok(response)
        });
    }

    {
        let indices = indices.clone();

        router.put("/:index", move |req: &mut Request| -> IronResult<Response> {
            // URL parameters
            let ref index_name = req.extensions.get::<Router>().unwrap().find("index").unwrap_or("");

            // Lock index array
            let mut indices = indices.lock().unwrap();

            // Load data from body
            let mut payload = String::new();
            req.body.read_to_string(&mut payload).unwrap();

            let data = if !payload.is_empty() {
                Some(match Json::from_str(&payload) {
                    Ok(data) => data,
                    Err(error) => {
                        // TODO: What specifically is bad about the JSON?
                        let mut response = Response::with((status::BadRequest, "{\"message\": \"Couldn't parse JSON\"}"));
                        response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
                        return Ok(response);
                    }
                })
            } else {
                None
            };

            // Create index
            indices.insert(index_name.clone().to_owned(), Index::new());

            let mut response = Response::with((status::Ok, "{\"acknowledged\": true}"));
            response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
            Ok(response)
        });
    }

    {
        let indices = indices.clone();

        router.put("/:index/_mapping/:mapping", move |req: &mut Request| -> IronResult<Response> {
            // URL parameters
            let index_name = req.extensions.get::<Router>().unwrap().find("index").unwrap_or("");
            let ref mapping_name = req.extensions.get::<Router>().unwrap().find("mapping").unwrap_or("");

            // Lock index array
            let mut indices = indices.lock().unwrap();

            // Find index
            let mut index = match indices.get_mut(index_name) {
                Some(index) => index,
                None => {
                    return Ok(index_not_found_response());
                }
            };

            // Load data from body
            let mut payload = String::new();
            req.body.read_to_string(&mut payload).unwrap();

            let data = if !payload.is_empty() {
                Some(match Json::from_str(&payload) {
                    Ok(data) => data,
                    Err(error) => {
                        // TODO: What specifically is bad about the JSON?
                        let mut response = Response::with((status::BadRequest, "{\"message\": \"Couldn't parse JSON\"}"));
                        response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
                        return Ok(response);
                    }
                })
            } else {
                None
            };

            // Insert mapping
            index.mappings.insert(mapping_name.clone().to_owned(), Mapping::new());

            let mut response = Response::with((status::Ok, "{\"acknowledged\": true}"));
            response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
            Ok(response)
        });
    }

    {
        let indices = indices.clone();

        router.post("/_bulk", move |req: &mut Request| -> IronResult<Response> {
            // Lock index array
            let mut indices = indices.lock().unwrap();

            // Load data from body
            let mut payload = String::new();
            req.body.read_to_string(&mut payload).unwrap();

            for payload_part in payload.split('\n') {
                let data = if !payload.is_empty() {
                    Some(Json::from_str(&payload))
                } else {
                    None
                };
            }

            let mut response = Response::with((status::Ok, "{\"acknowledged\": true}"));
            response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
            Ok(response)
        });
    }

    Iron::new(router).http("localhost:9200").unwrap();
}
<commit_msg>Added globals struct<commit_after>extern crate iron;
extern crate router;
extern crate rustc_serialize;

use std::io::Read;
use std::sync::{Arc, Mutex};
use std::collections::HashMap;

use iron::prelude::*;
use iron::status;
use router::Router;
use rustc_serialize::json::{self, Json};

#[derive(Debug)]
struct Document {
    data: Json,
}

impl Document {
    fn from_json(data: Json) -> Document {
        Document{
            data: data,
        }
    }
}

#[derive(Debug)]
enum Filter {
    Not(Box<Filter>),
    Or(Vec<Filter>),
    And(Vec<Filter>),
    Term(String, String),
}

impl Filter {
    fn matches(&self, doc: &Document) -> bool {
        match *self {
            Filter::Not(ref filter) => !filter.matches(doc),
            Filter::Or(ref filters) => {
                for filter in filters.iter() {
                    if (filter.matches(doc)) {
                        return true;
                    }
                }

                false
            },
            Filter::And(ref filters) => {
                for filter in filters.iter() {
                    if (!filter.matches(doc)) {
                        return false;
                    }
                }

                true
            },
            Filter::Term(ref field, ref value) => {
                let obj = doc.data.as_object().unwrap();

                if let Some(field_value) = obj.get(field) {
                    if let Json::String(ref field_value) = *field_value {
                        return field_value == value
                    }
                }

                false
            }
        }
    }
}

#[derive(Debug)]
struct Mapping {
    pub docs: HashMap<String, Document>,
}

impl Mapping {
    fn new() -> Mapping {
        Mapping{
            docs: HashMap::new(),
        }
    }
}

#[derive(Debug)]
struct Index {
    pub mappings: HashMap<String, Mapping>,
}

impl Index {
    fn new() -> Index {
        Index{
            mappings: HashMap::new(),
        }
    }
}

struct Globals {
    pub indices: Mutex<HashMap<String, Index>>,
}

impl Globals {
    fn new() -> Globals {
        Globals {
            indices: Mutex::new(HashMap::new())
        }
    }
}

fn index_not_found_response() -> Response {
    let mut response = Response::with((status::NotFound, "{\"message\": \"Index not found\"}"));
    response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
    return response;
}

fn main() {
    let glob = Arc::new(Globals::new());
    let mut wagtail_index = Index::new();
    wagtail_index.mappings.insert("wagtaildocs_document".to_owned(), Mapping::new());
    glob.indices.lock().unwrap().insert("wagtaildemo".to_owned(), wagtail_index);

    let f = Filter::Or(vec![
        Filter::Term("title".to_owned(), "test".to_owned()),
        Filter::Term("title".to_owned(), "foo".to_owned()),
    ]);

    println!("{:?}", f);

    let mut router = Router::new();

    router.get("/", |_: &mut Request| -> IronResult<Response> {
        Ok(Response::with((status::Ok, "Hello World!")))
    });

    {
        let glob = glob.clone();

        router.get("/:index/_count", move |req: &mut Request| -> IronResult<Response> {
            // URL parameters
            let index_name = req.extensions.get::<Router>().unwrap().find("index").unwrap_or("");

            // Lock index array
            let indices = glob.indices.lock().unwrap();

            // Find index
            let index = match indices.get(index_name) {
                Some(index) => index,
                None => {
                    return Ok(index_not_found_response());
                }
            };

            // Load query from body
            let mut payload = String::new();
            req.body.read_to_string(&mut payload).unwrap();

            let data = if !payload.is_empty() {
                Some(match Json::from_str(&payload) {
                    Ok(data) => data,
                    Err(error) => {
                        // TODO: What specifically is bad about the JSON?
                        let mut response = Response::with((status::BadRequest, "{\"message\": \"Couldn't parse JSON\"}"));
                        response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
                        return Ok(response);
                    }
                })
            } else {
                None
            };

            // TODO: Run query

            // Temporary count and return numbers
            let mut count = 0;
            for mapping in index.mappings.values() {
                count += mapping.docs.len();
            }

            let mut response = Response::with((status::Ok, format!("{{\"count\": {}}}", count)));
            response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
            Ok(response)
        });
    }

    {
        let glob = glob.clone();

        router.get("/:index/_search", move |req: &mut Request| -> IronResult<Response> {
            // URL parameters
            let index_name = req.extensions.get::<Router>().unwrap().find("index").unwrap_or("");

            // Lock index array
            let indices = glob.indices.lock().unwrap();

            // Find index
            let index = match indices.get(index_name) {
                Some(index) => index,
                None => {
                    return Ok(index_not_found_response());
                }
            };

            // Load query from body
            let mut payload = String::new();
            req.body.read_to_string(&mut payload).unwrap();

            let data = if !payload.is_empty() {
                Some(match Json::from_str(&payload) {
                    Ok(data) => data,
                    Err(error) => {
                        // TODO: What specifically is bad about the JSON?
                        let mut response = Response::with((status::BadRequest, "{\"message\": \"Couldn't parse JSON\"}"));
                        response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
                        return Ok(response);
                    }
                })
            } else {
                None
            };

            // TODO: Run query

            let mut response = Response::with((status::Ok, "{\"hits\": {\"total\": 0, \"hits\": []}}"));
            response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
            Ok(response)
        });
    }

    {
        let glob = glob.clone();

        router.get("/:index/:mapping/:doc", move |req: &mut Request| -> IronResult<Response> {
            // URL parameters
            let index_name = req.extensions.get::<Router>().unwrap().find("index").unwrap_or("");
            let mapping_name = req.extensions.get::<Router>().unwrap().find("mapping").unwrap_or("");
            let doc_id = req.extensions.get::<Router>().unwrap().find("doc").unwrap_or("");

            // Lock index array
            let indices = glob.indices.lock().unwrap();

            // Find index
            let index = match indices.get(index_name) {
                Some(index) => index,
                None => {
                    return Ok(index_not_found_response());
                }
            };

            // Find mapping
            let mapping = match index.mappings.get(mapping_name) {
                Some(mapping) => mapping,
                None => {
                    let mut response = Response::with((status::NotFound, "{\"message\": \"Mapping not found\"}"));
                    response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
                    return Ok(response);
                }
            };

            // Find document
            let doc = match mapping.docs.get(doc_id) {
                Some(doc) => doc,
                None => {
                    let mut response = Response::with((status::NotFound, "{\"message\": \"Document not found\"}"));
                    response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
                    return Ok(response);
                }
            };

            let mut response = Response::with((status::Ok, json::encode(&doc.data).unwrap()));
            response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
            Ok(response)
        });
    }

    {
        let glob = glob.clone();

        router.put("/:index/:mapping/:doc", move |req: &mut Request| -> IronResult<Response> {
            // URL parameters
            let index_name = req.extensions.get::<Router>().unwrap().find("index").unwrap_or("");
            let mapping_name = req.extensions.get::<Router>().unwrap().find("mapping").unwrap_or("");
            let ref doc_id = req.extensions.get::<Router>().unwrap().find("doc").unwrap_or("");

            // Lock index array
            let mut indices = glob.indices.lock().unwrap();

            // Find index
            let mut index = match indices.get_mut(index_name) {
                Some(index) => index,
                None => {
                    return Ok(index_not_found_response());
                }
            };

            // Find mapping
            let mut mapping = match index.mappings.get_mut(mapping_name) {
                Some(mapping) => mapping,
                None => {
                    let mut response = Response::with((status::NotFound, "{\"message\": \"Mapping not found\"}"));
                    response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
                    return Ok(response);
                }
            };

            // Load data from body
            let mut payload = String::new();
            req.body.read_to_string(&mut payload).unwrap();

            let data = if !payload.is_empty() {
                Some(match Json::from_str(&payload) {
                    Ok(data) => data,
                    Err(error) => {
                        // TODO: What specifically is bad about the JSON?
                        let mut response = Response::with((status::BadRequest, "{\"message\": \"Couldn't parse JSON\"}"));
                        response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
                        return Ok(response);
                    }
                })
            } else {
                None
            };

            // Create and insert document
            if let Some(data) = data {
                let doc = Document::from_json(data);
                println!("{:?}", f.matches(&doc));
                mapping.docs.insert(doc_id.clone().to_owned(), doc);
            }

            let mut response = Response::with((status::Ok, "{}"));
            response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
            Ok(response)
        });
    }

    {
        let glob = glob.clone();

        router.put("/:index", move |req: &mut Request| -> IronResult<Response> {
            // URL parameters
            let ref index_name = req.extensions.get::<Router>().unwrap().find("index").unwrap_or("");

            // Lock index array
            let mut indices = glob.indices.lock().unwrap();

            // Load data from body
            let mut payload = String::new();
            req.body.read_to_string(&mut payload).unwrap();

            let data = if !payload.is_empty() {
                Some(match Json::from_str(&payload) {
                    Ok(data) => data,
                    Err(error) => {
                        // TODO: What specifically is bad about the JSON?
                        let mut response = Response::with((status::BadRequest, "{\"message\": \"Couldn't parse JSON\"}"));
                        response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
                        return Ok(response);
                    }
                })
            } else {
                None
            };

            // Create index
            indices.insert(index_name.clone().to_owned(), Index::new());

            let mut response = Response::with((status::Ok, "{\"acknowledged\": true}"));
            response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
            Ok(response)
        });
    }

    {
        let glob = glob.clone();

        router.put("/:index/_mapping/:mapping", move |req: &mut Request| -> IronResult<Response> {
            // URL parameters
            let index_name = req.extensions.get::<Router>().unwrap().find("index").unwrap_or("");
            let ref mapping_name = req.extensions.get::<Router>().unwrap().find("mapping").unwrap_or("");

            // Lock index array
            let mut indices = glob.indices.lock().unwrap();

            // Find index
            let mut index = match indices.get_mut(index_name) {
                Some(index) => index,
                None => {
                    return Ok(index_not_found_response());
                }
            };

            // Load data from body
            let mut payload = String::new();
            req.body.read_to_string(&mut payload).unwrap();

            let data = if !payload.is_empty() {
                Some(match Json::from_str(&payload) {
                    Ok(data) => data,
                    Err(error) => {
                        // TODO: What specifically is bad about the JSON?
                        let mut response = Response::with((status::BadRequest, "{\"message\": \"Couldn't parse JSON\"}"));
                        response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
                        return Ok(response);
                    }
                })
            } else {
                None
            };

            // Insert mapping
            index.mappings.insert(mapping_name.clone().to_owned(), Mapping::new());

            let mut response = Response::with((status::Ok, "{\"acknowledged\": true}"));
            response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
            Ok(response)
        });
    }

    {
        let glob = glob.clone();

        router.post("/_bulk", move |req: &mut Request| -> IronResult<Response> {
            // Lock index array
            let mut indices = glob.indices.lock().unwrap();

            // Load data from body
            let mut payload = String::new();
            req.body.read_to_string(&mut payload).unwrap();

            for payload_part in payload.split('\n') {
                let data = if !payload.is_empty() {
                    Some(Json::from_str(&payload))
                } else {
                    None
                };
            }

            let mut response = Response::with((status::Ok, "{\"acknowledged\": true}"));
            response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
            Ok(response)
        });
    }

    Iron::new(router).http("localhost:9200").unwrap();
}
<|endoftext|>


<commit_before>// The MIT License (MIT)
//
// Copyright (c) 2015 Johan Johansson
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

// TODO: TESTS!
// TODO: Color shift
// TODO: Support capturing a virtual desktop, multiple monotors together
// TODO: Optional solid color or animation when there's no signal

extern crate serde;
#[macro_use]
extern crate serde_derive;
extern crate serde_json;
extern crate clock_ticks;
extern crate serial;
extern crate captrs;
extern crate simd;

#[cfg(feature = "cpuprofiler")]
extern crate cpuprofiler;

use captrs::{Capturer, CaptureError};
use capture::ImageAnalyzer;
use color::{Rgb8, RgbTransformer, HSVTransformer, Color};
use config::parse_led_indices;
use serial::SerialPort;
use std::{thread, time};
use std::cmp::{max, Ordering};

mod config;
mod color;
mod capture;

type SmoothFn = fn(Rgb8, Rgb8, f32) -> Rgb8;

/// Returns smallest of `a` and `b` if there is one, else returns `expect`
fn partial_min<T: PartialOrd>(a: T, b: T, expect: T) -> T {
    match a.partial_cmp(&b) {
        Some(Ordering::Less) |
        Some(Ordering::Equal) => a,
        Some(Ordering::Greater) => b,
        None => expect,
    }
}
/// Returns greatest of `a` and `b` if there is one, else returns `expect`
fn partial_max<T: PartialOrd>(a: T, b: T, expect: T) -> T {
    match a.partial_cmp(&b) {
        Some(Ordering::Greater) |
        Some(Ordering::Equal) => a,
        Some(Ordering::Less) => b,
        None => expect,
    }
}

/// A special header is expected by the corresponding LED streaming code running on the Arduino.
/// This only needs to be initialized once since the number of LEDs remains constant.
fn new_pixel_buf_header(n_leds: u16) -> [u8; 6] {
    // In the two below, not sure why the -1 in `(n_leds - 1)` is needed,
    // but that's how LEDstream on the Arduino expects it
    let count_high = ((n_leds - 1) >> 8) as u8; // LED count high byte
    let count_low = ((n_leds - 1) & 0xff) as u8; // LED count low byte

    ['A' as u8,
     'd' as u8,
     'a' as u8,
     count_high,
     count_low,
     count_high ^ count_low ^ 0x55 /* Checksum */]
}

/// A timer to track time passed between refreshes. Can for example be used to limit frame rate.
struct FrameTimer {
    before: f64,
    last_frame_dt: f64,
}
impl FrameTimer {
    fn new() -> FrameTimer {
        FrameTimer {
            before: clock_ticks::precise_time_s(),
            last_frame_dt: 0.0,
        }
    }

    /// For how long the previous frame lasted
    fn last_frame_dt(&self) -> f64 {
        self.last_frame_dt
    }

    /// Time passed since last tick
    fn dt_to_now(&mut self) -> f64 {
        let now = clock_ticks::precise_time_s();
        let dt = now - self.before;
        if dt >= 0.0 {
            dt
        } else {
            self.before = now;
            0.0
        }
    }

    /// An update/frame/refresh has occured; take the time.
    fn tick(&mut self) {
        let now = clock_ticks::precise_time_s();
        self.last_frame_dt = partial_max(now - self.before, 0.0, 0.0);
        self.before = now;
    }
}

/// Writes color data via serial to LEDstream compatible device
struct ColorWriter {
    con: serial::SystemPort,
    header: [u8; 6],
}

impl ColorWriter {
    /// Configure serial writing given a serial port, baud rate, and header to write before
    /// each data write
    fn new(port: &str, baud_rate: u32, header: [u8; 6]) -> Self {
        let mut serial_con = serial::open(port).unwrap();

        let baud_rate = serial::BaudRate::from_speed(baud_rate as usize);

        serial_con.reconfigure(&|cfg| cfg.set_baud_rate(baud_rate))
                  .unwrap();

        ColorWriter { con: serial_con, header: header }
    }

    /// Write a buffer of color data to the LEDstream device
    fn write_colors(&mut self, color_data: &[Rgb8]) {
        use std::io::Write;

        let color_bytes = color::rgbs_as_bytes(color_data);

        match self.con.write(&self.header) {
            Ok(hn) if hn == self.header.len() => {
                match self.con.write(color_bytes) {
                    Ok(bn) if bn == color_bytes.len() => (),
                    Ok(bn) => {
                        println!("Failed to write all bytes of RGB data. Wrote {} of {}",
                                 bn,
                                 color_data.len())
                    }
                    Err(e) => println!("Failed to write RGB data, {}", e),
                }
            }
            Ok(_) => println!("Failed to write all bytes in header"),
            Err(e) => println!("Failed to write header, {}", e),
        }
    }
}

/// Update the colors to output by analyzing the captured frame for each led
fn update_out_color_data(out_pixels: &mut [Rgb8],
                         frame_analyzer: &ImageAnalyzer,
                         leds: &[config::Region],
                         leds_transformers: &[Vec<(Option<RgbTransformer>,
                                                   Option<&HSVTransformer>)>],
                         smooth: SmoothFn,
                         smooth_factor: f32) {
    for (i, &led) in leds.iter().enumerate() {
        let avg_color = frame_analyzer.average_color(led);

        let to_pixel = leds_transformers[i]
            .iter()
            .map(|&(ref opt_rgb, ref opt_hsv)| (opt_rgb.as_ref(), opt_hsv.as_ref()))
            .fold(Color::RGB(avg_color), |acc_color, transformers| {
                match transformers {
                    (Some(rgb_tr), Some(hsv_tr)) => {
                        Color::HSV(hsv_tr.transform(rgb_tr.transform(acc_color.into_rgb())
                                                          .to_hsv()))
                    }
                    (Some(rgb_tr), _) => Color::RGB(rgb_tr.transform(acc_color.into_rgb())),
                    (_, Some(hsv_tr)) => Color::HSV(hsv_tr.transform(acc_color.into_hsv())),
                    _ => acc_color,
                }
            })
            .into_rgb();

        out_pixels[i] = smooth(out_pixels[i], to_pixel, smooth_factor);
    }
}

/// Specifies how to loop over the main body
#[cfg(not(feature = "cpuprofiler"))]
fn main_loop<F: FnMut() -> bool>(mut body: F) {
    loop {
        if !body() {
            break;
        }
    }
}

/// Loop over the main body a fixed number of times and track performance with a profiler
#[cfg(feature = "cpuprofiler")]
fn main_loop<F: FnMut() -> bool>(mut body: F) {
    cpuprofiler::PROFILER.lock()
                         .unwrap()
                         .start("./prof.profile")
                         .unwrap();

    for _ in 0..1000 {
        if !body() {
            break;
        }
    }

    cpuprofiler::PROFILER.lock().unwrap().stop().unwrap();
}

fn main() {
    let config = config::parse_config();

    let leds: &[_] = &config.leds;

    let mut led_transformers_list: Vec<_> = vec![Vec::with_capacity(1); leds.len()];

    // Add color transforms from config to each led in matching vec
    for transform_conf in config.color.transform.iter() {
        let hsv_transformer =
            if !transform_conf.hsv.is_default() { Some(&transform_conf.hsv) } else { None };

        let rgb_transformer = if !(transform_conf.red.is_default() &&
                                   transform_conf.green.is_default() &&
                                   transform_conf.red.is_default()) {
            Some(RgbTransformer {
                     r: transform_conf.red.clone(),
                     g: transform_conf.green.clone(),
                     b: transform_conf.blue.clone(),
                 })
        } else {
            None
        };

        for range in parse_led_indices(&transform_conf.leds, leds.len()).iter() {
            for transformers in led_transformers_list[range.clone()].iter_mut() {
                transformers.push((rgb_transformer.clone(), hsv_transformer));
            }
        }
    }

    // Header to write before led data
    let out_header = new_pixel_buf_header(leds.len() as u16);

    let mut color_writer = ColorWriter::new(&config.device.output, config.device.rate, out_header);

    // Skeleton for the output led pixel buffer to write to arduino
    let mut out_pixels = vec![Rgb8 { r: 0, g: 0, b: 0 }; leds.len()];

    let mut capturer = Capturer::new(config.device.input as usize).unwrap();

    let capture_frame_interval = 1.0 / config.framegrabber.frequency_Hz;

    // Function to use when smoothing led colors
    let smooth = match config.color.smoothing.type_.as_ref() {
        "linear" => color::linear_smooth as SmoothFn,
        _ => color::no_smooth as SmoothFn,
    };

    // max w/ 1 to avoid future divide by zero
    let smooth_time_const = max(config.color.smoothing.time_ms, 1) as f64 / 1000.0;

    let led_refresh_interval = 1.0 / config.color.smoothing.update_frequency;

    println!("Helion - An LED streamer\nNumber of LEDs: {}\nResize resolution: {} x {}\nCapture \
              rate: {} fps\nLED refresh rate: {} hz\nCapture output: {}\nSerial port: {}",
             leds.len(),
             config.framegrabber.width,
             config.framegrabber.height,
             config.framegrabber.frequency_Hz,
             1.0 / led_refresh_interval,
             config.device.input,
             config.device.output);

    let mut capture_timer = FrameTimer::new();
    let mut led_refresh_timer = FrameTimer::new();

    main_loop(|| {
        led_refresh_timer.tick();

        // Don't capture new frame if going faster than frame limit,
        // but still proceed to smooth leds
        if capture_timer.dt_to_now() > capture_frame_interval {
            // If something goes wrong, last frame is reused

            match capturer.capture_store_frame() {
                Ok(_) => (),
                Err(CaptureError::Timeout) => {
                    println!("timeout");
                    thread::sleep(time::Duration::from_millis(100));
                    return true;
                }
                Err(e) => {
                    println!("Capture error: {:?}", e);
                    thread::sleep(time::Duration::from_millis(2_000));

                    loop {
                        match Capturer::new(config.device.input as usize) {
                            Ok(c) => {
                                capturer = c;
                                break;
                            }
                            Err(e1) => {
                                println!("Error creating capturer: {:?}", e1);
                                thread::sleep(time::Duration::from_millis(1_000));
                            }
                        }
                    }
                }
            }

            capture_timer.tick();
        }

        if let Some(frame) = capturer.get_stored_frame() {
            let (w, h) = capturer.geometry();
            let frame_analyzer = ImageAnalyzer::new(frame,
                                                    w as usize,
                                                    h as usize,
                                                    config.framegrabber.width,
                                                    config.framegrabber.height);

            let smooth_factor = (led_refresh_timer.last_frame_dt() / smooth_time_const) as f32;

            update_out_color_data(&mut out_pixels,
                                  &frame_analyzer,
                                  leds,
                                  &led_transformers_list,
                                  smooth,
                                  smooth_factor);

            color_writer.write_colors(&out_pixels)
        }

        let time_left = led_refresh_interval - led_refresh_timer.dt_to_now();
        if time_left > 0.0 {
            let ms = if time_left > 0.0 { time_left * 1_000.0 } else { 0.0 };

            thread::sleep(time::Duration::from_millis(ms as u64));
        }
        true
    })
}
<commit_msg>Retry opening serial dev on conn loss & Fix serial conf<commit_after>// The MIT License (MIT)
//
// Copyright (c) 2015 Johan Johansson
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

// TODO: TESTS!
// TODO: Color shift
// TODO: Support capturing a virtual desktop, multiple monotors together
// TODO: Optional solid color or animation when there's no signal

extern crate serde;
#[macro_use]
extern crate serde_derive;
extern crate serde_json;
extern crate clock_ticks;
extern crate serial;
extern crate captrs;
extern crate simd;

#[cfg(feature = "cpuprofiler")]
extern crate cpuprofiler;

use captrs::{Capturer, CaptureError};
use capture::ImageAnalyzer;
use color::{Rgb8, RgbTransformer, HSVTransformer, Color};
use config::parse_led_indices;
use serial::SerialPort;
use std::{thread, time};
use std::cmp::{max, Ordering};

mod config;
mod color;
mod capture;

type SmoothFn = fn(Rgb8, Rgb8, f32) -> Rgb8;

/// Returns smallest of `a` and `b` if there is one, else returns `expect`
fn partial_min<T: PartialOrd>(a: T, b: T, expect: T) -> T {
    match a.partial_cmp(&b) {
        Some(Ordering::Less) |
        Some(Ordering::Equal) => a,
        Some(Ordering::Greater) => b,
        None => expect,
    }
}
/// Returns greatest of `a` and `b` if there is one, else returns `expect`
fn partial_max<T: PartialOrd>(a: T, b: T, expect: T) -> T {
    match a.partial_cmp(&b) {
        Some(Ordering::Greater) |
        Some(Ordering::Equal) => a,
        Some(Ordering::Less) => b,
        None => expect,
    }
}

/// A special header is expected by the corresponding LED streaming code running on the Arduino.
/// This only needs to be initialized once since the number of LEDs remains constant.
fn new_pixel_buf_header(n_leds: u16) -> [u8; 6] {
    // In the two below, not sure why the -1 in `(n_leds - 1)` is needed,
    // but that's how LEDstream on the Arduino expects it
    let count_high = ((n_leds - 1) >> 8) as u8; // LED count high byte
    let count_low = ((n_leds - 1) & 0xff) as u8; // LED count low byte

    ['A' as u8,
     'd' as u8,
     'a' as u8,
     count_high,
     count_low,
     count_high ^ count_low ^ 0x55 /* Checksum */]
}

/// A timer to track time passed between refreshes. Can for example be used to limit frame rate.
struct FrameTimer {
    before: f64,
    last_frame_dt: f64,
}
impl FrameTimer {
    fn new() -> FrameTimer {
        FrameTimer {
            before: clock_ticks::precise_time_s(),
            last_frame_dt: 0.0,
        }
    }

    /// For how long the previous frame lasted
    fn last_frame_dt(&self) -> f64 {
        self.last_frame_dt
    }

    /// Time passed since last tick
    fn dt_to_now(&mut self) -> f64 {
        let now = clock_ticks::precise_time_s();
        let dt = now - self.before;
        if dt >= 0.0 {
            dt
        } else {
            self.before = now;
            0.0
        }
    }

    /// An update/frame/refresh has occured; take the time.
    fn tick(&mut self) {
        let now = clock_ticks::precise_time_s();
        self.last_frame_dt = partial_max(now - self.before, 0.0, 0.0);
        self.before = now;
    }
}

/// Writes color data via serial to LEDstream compatible device
struct ColorWriter {
    port: String,
    baud_rate: serial::BaudRate,
    con: serial::SystemPort,
    header: [u8; 6],
}

impl ColorWriter {
    /// Configure serial writing given a serial port, baud rate, and header to write before
    /// each data write
    fn new(port: &str, baud_rate: u32, header: [u8; 6]) -> Self {
        let serial_baud_rate = serial::BaudRate::from_speed(baud_rate as usize);

        let serial_con = ColorWriter::connect(port, serial_baud_rate).unwrap();

        ColorWriter {
            port: port.to_string(),
            baud_rate: serial_baud_rate,
            con: serial_con,
            header: header,
        }
    }

    fn connect(port: &str,
               baud_rate: serial::BaudRate)
               -> Result<serial::SystemPort, serial::Error> {
        serial::open(port).and_then(|mut con| {
            con.reconfigure(&|settings| {
                settings.set_char_size(serial::CharSize::Bits8);
                settings.set_baud_rate(baud_rate)
            })
                .map(|()| con)
        })
    }

    /// Write a buffer of color data to the LEDstream device
    fn write_colors(&mut self, color_data: &[Rgb8]) {
        use std::io::Write;

        let color_bytes = color::rgbs_as_bytes(color_data);

        match self.con.write(&self.header) {
            Ok(hn) if hn == self.header.len() => {
                match self.con.write(color_bytes) {
                    Ok(bn) if bn == color_bytes.len() => (),
                    Ok(bn) => {
                        println!("Failed to write all bytes of RGB data. Wrote {} of {}",
                                 bn,
                                 color_data.len())
                    }
                    Err(e) => println!("Failed to write RGB data, {}", e),
                }
            }
            Ok(_) => println!("Failed to write all bytes in header"),
            Err(e) => {
                println!("Failed to write header, {}", e);
                // Probably, the serial device was temporarily disconnected or rebooted.
                // Try to reconnect forever
                loop {
                    print!("Reopening serial port... ");
                    match ColorWriter::connect(&self.port, self.baud_rate) {
                        Ok(con) => {
                            self.con = con;
                            println!("success!");
                            break;
                        }
                        Err(e) => {
                            println!("{}", e);
                            thread::sleep(time::Duration::from_millis(1000));
                        }
                    }
                }
            }
        }
    }
}

/// Update the colors to output by analyzing the captured frame for each led
fn update_out_color_data(out_pixels: &mut [Rgb8],
                         frame_analyzer: &ImageAnalyzer,
                         leds: &[config::Region],
                         leds_transformers: &[Vec<(Option<RgbTransformer>,
                                                   Option<&HSVTransformer>)>],
                         smooth: SmoothFn,
                         smooth_factor: f32) {
    for (i, &led) in leds.iter().enumerate() {
        let avg_color = frame_analyzer.average_color(led);

        let to_pixel = leds_transformers[i]
            .iter()
            .map(|&(ref opt_rgb, ref opt_hsv)| (opt_rgb.as_ref(), opt_hsv.as_ref()))
            .fold(Color::RGB(avg_color), |acc_color, transformers| {
                match transformers {
                    (Some(rgb_tr), Some(hsv_tr)) => {
                        Color::HSV(hsv_tr.transform(rgb_tr.transform(acc_color.into_rgb())
                                                          .to_hsv()))
                    }
                    (Some(rgb_tr), _) => Color::RGB(rgb_tr.transform(acc_color.into_rgb())),
                    (_, Some(hsv_tr)) => Color::HSV(hsv_tr.transform(acc_color.into_hsv())),
                    _ => acc_color,
                }
            })
            .into_rgb();

        out_pixels[i] = smooth(out_pixels[i], to_pixel, smooth_factor);
    }
}

/// Specifies how to loop over the main body
#[cfg(not(feature = "cpuprofiler"))]
fn main_loop<F: FnMut() -> bool>(mut body: F) {
    loop {
        if !body() {
            break;
        }
    }
}

/// Loop over the main body a fixed number of times and track performance with a profiler
#[cfg(feature = "cpuprofiler")]
fn main_loop<F: FnMut() -> bool>(mut body: F) {
    cpuprofiler::PROFILER.lock()
                         .unwrap()
                         .start("./prof.profile")
                         .unwrap();

    for _ in 0..1000 {
        if !body() {
            break;
        }
    }

    cpuprofiler::PROFILER.lock().unwrap().stop().unwrap();
}

fn main() {
    let config = config::parse_config();

    let leds: &[_] = &config.leds;

    let mut led_transformers_list: Vec<_> = vec![Vec::with_capacity(1); leds.len()];

    // Add color transforms from config to each led in matching vec
    for transform_conf in config.color.transform.iter() {
        let hsv_transformer =
            if !transform_conf.hsv.is_default() { Some(&transform_conf.hsv) } else { None };

        let rgb_transformer = if !(transform_conf.red.is_default() &&
                                   transform_conf.green.is_default() &&
                                   transform_conf.red.is_default()) {
            Some(RgbTransformer {
                     r: transform_conf.red.clone(),
                     g: transform_conf.green.clone(),
                     b: transform_conf.blue.clone(),
                 })
        } else {
            None
        };

        for range in parse_led_indices(&transform_conf.leds, leds.len()).iter() {
            for transformers in led_transformers_list[range.clone()].iter_mut() {
                transformers.push((rgb_transformer.clone(), hsv_transformer));
            }
        }
    }

    // Header to write before led data
    let out_header = new_pixel_buf_header(leds.len() as u16);

    let mut color_writer = ColorWriter::new(&config.device.output, config.device.rate, out_header);

    // Skeleton for the output led pixel buffer to write to arduino
    let mut out_pixels = vec![Rgb8 { r: 0, g: 0, b: 0 }; leds.len()];

    let mut capturer = Capturer::new(config.device.input as usize).unwrap();

    let capture_frame_interval = 1.0 / config.framegrabber.frequency_Hz;

    // Function to use when smoothing led colors
    let smooth = match config.color.smoothing.type_.as_ref() {
        "linear" => color::linear_smooth as SmoothFn,
        _ => color::no_smooth as SmoothFn,
    };

    // max w/ 1 to avoid future divide by zero
    let smooth_time_const = max(config.color.smoothing.time_ms, 1) as f64 / 1000.0;

    let led_refresh_interval = 1.0 / config.color.smoothing.update_frequency;

    println!("Helion - An LED streamer\nNumber of LEDs: {}\nResize resolution: {} x {}\nCapture \
              rate: {} fps\nLED refresh rate: {} hz\nCapture output: {}\nSerial port: {}",
             leds.len(),
             config.framegrabber.width,
             config.framegrabber.height,
             config.framegrabber.frequency_Hz,
             1.0 / led_refresh_interval,
             config.device.input,
             config.device.output);

    let mut capture_timer = FrameTimer::new();
    let mut led_refresh_timer = FrameTimer::new();

    main_loop(|| {
        led_refresh_timer.tick();

        // Don't capture new frame if going faster than frame limit,
        // but still proceed to smooth leds
        if capture_timer.dt_to_now() > capture_frame_interval {
            // If something goes wrong, last frame is reused

            match capturer.capture_store_frame() {
                Ok(_) => (),
                Err(CaptureError::Timeout) => {
                    println!("timeout");
                    thread::sleep(time::Duration::from_millis(100));
                    return true;
                }
                Err(e) => {
                    println!("Capture error: {:?}", e);
                    thread::sleep(time::Duration::from_millis(2_000));

                    loop {
                        match Capturer::new(config.device.input as usize) {
                            Ok(c) => {
                                capturer = c;
                                break;
                            }
                            Err(e1) => {
                                println!("Error creating capturer: {:?}", e1);
                                thread::sleep(time::Duration::from_millis(1_000));
                            }
                        }
                    }
                }
            }

            capture_timer.tick();
        }

        if let Some(frame) = capturer.get_stored_frame() {
            let (w, h) = capturer.geometry();
            let frame_analyzer = ImageAnalyzer::new(frame,
                                                    w as usize,
                                                    h as usize,
                                                    config.framegrabber.width,
                                                    config.framegrabber.height);

            let smooth_factor = (led_refresh_timer.last_frame_dt() / smooth_time_const) as f32;

            update_out_color_data(&mut out_pixels,
                                  &frame_analyzer,
                                  leds,
                                  &led_transformers_list,
                                  smooth,
                                  smooth_factor);

            color_writer.write_colors(&out_pixels)
        }

        let time_left = led_refresh_interval - led_refresh_timer.dt_to_now();
        if time_left > 0.0 {
            let ms = if time_left > 0.0 { time_left * 1_000.0 } else { 0.0 };

            thread::sleep(time::Duration::from_millis(ms as u64));
        }
        true
    })
}
<|endoftext|>


<commit_before>extern crate lancelot;
extern crate log;

use failure::{Error, Fail};
use log::{error, info, trace};
use std::env;
use std::process;

use lancelot::arch::*;
use lancelot::workspace::Workspace;

#[derive(Debug, Fail)]
pub enum MainError {
    #[fail(display = "foo")]
    Foo,
}

pub struct Config {
    pub filename: String,
}

impl Config {
    pub fn from_args(args: env::Args) -> Result<Config, &'static str> {
        let args: Vec<String> = args.collect();

        if args.len() < 2 {
            return Err("not enough arguments");
        }

        let filename = args[1].clone();
        trace!("config: parsed filename: {:?}", filename);

        Ok(Config { filename })
    }
}

pub fn setup_logging(_args: &Config) {
    simplelog::TermLogger::init(simplelog::LevelFilter::Info, simplelog::Config::default())
        .expect("failed to setup logging");
}

pub fn run(args: &Config) -> Result<(), Error> {
    info!("filename: {:?}", args.filename);

    let _ = Workspace::<Arch32>::from_file(&args.filename)?.load()?;

    Ok(())
}

fn main() {
    let args = Config::from_args(env::args()).unwrap_or_else(|err| {
        eprintln!("error parsing arguments: {}", err);
        process::exit(1);
    });

    setup_logging(&args);

    if let Err(e) = run(&args) {
        error!("{:?}", e)
    }
}
<commit_msg>main: accept mode via cli arg<commit_after>extern crate lancelot;
extern crate log;

use failure::{Error, Fail};
use log::{error, info, trace};
use std::env;
use std::process;

use lancelot::arch::*;
use lancelot::workspace::Workspace;

// TODO: removeme
use goblin::{Object};

#[derive(Debug, Fail)]
pub enum MainError {
    #[fail(display = "foo")]
    Foo,
}

pub struct Config {
    pub mode: u8,
    pub filename: String,
}

impl Config {
    pub fn from_args(args: env::Args) -> Result<Config, &'static str> {
        let args: Vec<String> = args.collect();

        if args.len() < 3 {
            return Err("not enough arguments: provide `lancelot.exe 32|64 /path/to/input`");
        }

        let mode = match args[1].as_ref() {
            "32" => 32,
            "64" => 64,
            _ => return Err("invalid mode, pick one of `32` or `64`"),
        };

        let filename = args[2].clone();
        trace!("config: parsed filename: {:?}", filename);

        Ok(Config { mode, filename })
    }
}

pub fn setup_logging(_args: &Config) {
    simplelog::TermLogger::init(simplelog::LevelFilter::Info, simplelog::Config::default())
        .expect("failed to setup logging");
}

pub fn run(args: &Config) -> Result<(), Error> {
    info!("filename: {:?}", args.filename);

    if args.mode == 32 {
        let mut ws = Workspace::<Arch32>::from_file(&args.filename)?.load()?;
    } else if args.mode == 64 {
        let mut ws = Workspace::<Arch64>::from_file(&args.filename)?.load()?;

        if ws.loader.get_name() == "Windows/64/PE" {
            if let Ok(Object::PE(pe)) = Object::parse(&ws.buf) {
                let entry = pe.entry;
                let exports: Vec<usize> = pe.exports.iter().map(|exp| exp.offset).collect();

                info!("PE entry: {:#x}", entry);
                ws.make_insn(entry as i64);
                for export in exports.iter() {
                    info!("export: {:#x}", export);
                    ws.make_insn(*export as i64);
                }
                ws.analyze();
            }
        }
    }

    Ok(())
}

fn main() {
    let args = Config::from_args(env::args()).unwrap_or_else(|err| {
        eprintln!("error parsing arguments: {}", err);
        process::exit(1);
    });

    setup_logging(&args);

    if let Err(e) = run(&args) {
        error!("{:?}", e)
    }
}
<|endoftext|>


<commit_before>fn main() {

    for i in 1..100 {
        println!(
            "{}",
            if div_by_fifteen(i) { "FizzBuzz".to_string() }
            else if div_by_three(i) { "Fizz".to_string() }
            else if div_by_five(i) { "Buzz".to_string() }
            else { i.to_string() });
    }

}

fn div_by_three(num: u64) -> bool {
    (num % 3) == 0
}

fn div_by_five(num: u64) -> bool {
    (num % 5) == 0
}

fn div_by_fifteen(num: u64) -> bool {
    div_by_five(num) && div_by_three(num)
}

#[cfg(test)]
mod test {

    use super::div_by_three;

    #[test]
    fn test_div_by_three_returns_false() {

        assert_eq!(false, div_by_three(1));
        assert_eq!(false, div_by_three(2));

        assert_eq!(false, div_by_three(4));
        assert_eq!(false, div_by_three(5));
    }

    #[test]
    fn test_div_by_three_returns_true() {

        for i in 1..300 {
            assert_eq!(true, div_by_three(3 * i));
        }
    }
}
<commit_msg>Add More Tests<commit_after>fn main() {

    for i in 1..100 {
        println!(
            "{}",
            if div_by_fifteen(i) { "FizzBuzz".to_string() }
            else if div_by_three(i) { "Fizz".to_string() }
            else if div_by_five(i) { "Buzz".to_string() }
            else { i.to_string() });
    }

}

fn div_by_three(num: u64) -> bool {
    (num % 3) == 0
}

fn div_by_five(num: u64) -> bool {
    (num % 5) == 0
}

fn div_by_fifteen(num: u64) -> bool {
    div_by_five(num) && div_by_three(num)
}

#[cfg(test)]
mod test {

    use super::div_by_three;
    use super::div_by_five;
    use super::div_by_fifteen;

    #[test]
    fn test_div_by_three_returns_false() {

        assert_eq!(false, div_by_three(1));
        assert_eq!(false, div_by_three(2));

        assert_eq!(false, div_by_three(4));
        assert_eq!(false, div_by_three(5));
    }

    #[test]
    fn test_div_by_three_returns_true() {

        for i in 1..300 {
            assert_eq!(true, div_by_three(3 * i));
        }
    }

    #[test]
    fn test_div_by_five_returns_false() {

        assert_eq!(false, div_by_five(3));
        assert_eq!(false, div_by_five(4));
        assert_eq!(false, div_by_five(1));
        assert_eq!(false, div_by_five(11));
        assert_eq!(false, div_by_five(19));
    }

    #[test]
    fn test_div_by_five_returns_true() {

        for i in 1..100 {
            assert_eq!(true, div_by_five(5 * i));
        }
    }

    #[test]
    fn test_div_by_fifteen_returns_false() {

        assert_eq!(false, div_by_fifteen(5));
        assert_eq!(false, div_by_fifteen(3));

        assert_eq!(false, div_by_fifteen(14));
        assert_eq!(false, div_by_fifteen(31));
    }

    #[test]
    fn test_div_by_fifteen_returns_true() {

        for i in 1..300 {
            assert_eq!(true, div_by_fifteen(15 * i));
        }
    }
}
<|endoftext|>