<commit_before>extern crate regex;
extern crate getopts;
#[macro_use]
extern crate lazy_static;
use std::collections::HashMap;
use std::{env, fs, io};
use std::io::prelude::*;
mod options;
use options::*;
use std::process;
use std::collections::hash_map::Entry::{Occupied, Vacant};

macro_rules! regex { ($re:expr) => { ::regex::Regex::new($re).unwrap() } }
macro_rules! println_stderr(
    ($($arg:tt)*) => (
        match writeln!(&mut ::std::io::stderr(), $($arg)* ) {
            Ok(_) => {},
            Err(x) => panic!("Unable to write to stderr: {}", x),
        }
        )
    );

const VERSION: &'static str = env!("CARGO_PKG_VERSION");
const KEY_FILE: &'static str = "KEY_FILE";
const ENV: &'static str = "ENV";
const MISSING: &'static str = "Missing properties";

fn replace_var(line: String,
               keys: &Option<KeysAndSources>,
               props_first: bool,
               used_keys: &mut HashMap<String, Vec<(String, String)>>,
               result: &mut Vec<String>) {
    let re = regex!(r"(?P<full>\$\{\s*(?P<var>[^\}]*)\s*\})");
    let mut str_line: String = line.clone();

    for cap in re.captures_iter(&*str_line.clone()) {
        cap.name("var")
           .map(|v| {
               let env_prop = env_or_prop(v.trim(), props_first, keys)
                                  .map(|v2| {
                                      str_line = re.replace(&*str_line,
                                                            &*(v2.1.value));
                                      (v2.0, v2.1)
                                  });

               env_prop.and_then(|v2| {
                   used_keys.get_mut(&v2.1.source)
                            .and_then(|r| {
                                Some(r.push((v2.0.clone(), v2.1.value.clone())))
                            })
                            .or_else(|| {
                                used_keys.insert(v2.1.source.clone(),
                                                 vec![(v2.0.clone(),
                                                       v2.1.value.clone())]);
                                Some(())
                            })
               })
           });
    }
    result.push(str_line.to_string());
}

fn env_or_prop(key: &str,
               props_first: bool,
               keys: &Option<KeysAndSources>)
               -> Option<(String, ValueAndSource)> {
    match keys {
        &Some(ref k) => {
            let dict = &(k.dictionary);
            let prop = dict.get(key).and_then(|v| Some(v.clone()));
            let env = env::var_os(key)
                          .map(|v2| v2.clone().into_string().ok().unwrap());
            match (prop, props_first) {
                (Some(prop), true) => {
                    Some((key.to_string(),
                          ValueAndSource {
                        value: prop.value.clone(),
                        source: prop.source.clone(),
                    }))
                }
                (Some(prop), false) => {
                    match env {
                        Some(e) => {
                            println!("[WARNING]: key \"${{{}}}\" from the \
                                      file \"{}\" was overridden by an \
                                      environment variable with the same \
                                      name.",
                                     key.to_string(),
                                     prop.source);
                            Some((key.to_string(),
                                  ValueAndSource {
                                value: e.to_string(),
                                source: ENV.to_string(),
                            }))
                        }
                        None => {
                            Some((key.to_string(),
                                  ValueAndSource {
                                value: prop.value.clone(),
                                source: prop.source.clone(),
                            }))
                        }
                    }
                }
                (None, true) => {
                    match env {
                        Some(e) => {
                            Some((key.to_string(),
                                  ValueAndSource {
                                value: e.to_string(),
                                source: ENV.to_string(),
                            }))
                        }
                        None => {
                            Some((key.to_string(),
                                  ValueAndSource {
                                value: key.to_string(),
                                source: MISSING.to_string(),
                            }))
                        }
                    }
                }
                (None, false) => {
                    match env {
                        Some(e) => {
                            Some((key.to_string(),
                                  ValueAndSource {
                                value: e.to_string(),
                                source: ENV.to_string(),
                            }))
                        }
                        None => {
                            Some((key.to_string(),
                                  ValueAndSource {
                                value: key.to_string(),
                                source: MISSING.to_string(),
                            }))
                        }
                    }
                }
            }
        }
        &None => {
            match env::var_os(key)
                      .map(|v| v.into_string().ok())
                      .unwrap_or(None) {
                Some(e) => {
                    Some((key.to_string(),
                          ValueAndSource {
                        value: e.to_string(),
                        source: ENV.to_string(),
                    }))
                }
                None => {
                    Some((key.to_string(),
                          ValueAndSource {
                        value: key.to_string(),
                        source: MISSING.to_string(),
                    }))
                }
            }
        }
    }
}

fn process(file: fs::File,
           used_keys: &mut HashMap<String, Vec<(String, String)>>,
           dict: &Option<KeysAndSources>,
           props_first: bool)
           -> Vec<String> {
    let buf_reader = io::BufReader::new(&file);
    let mut result: Vec<String> = vec![];

    for line in buf_reader.lines() {
        replace_var(line.unwrap(), &dict, props_first, used_keys, &mut result);
    }

    result
}

fn extract_keys(file: fs::File) -> HashMap<String, String> {
    let key_regex = regex!(r"^\s*(?P<key>\S*)\s*=\s*(?P<val>.*)");
    let buf_reader = io::BufReader::new(&file);
    let mut keys = HashMap::new();

    for line in buf_reader.lines().enumerate() {
        for cap in key_regex.captures_iter(&*line.1
                                                 .ok()
                                                 .expect("Could not read \
                                                          line")) {
            let key = cap.name("key")
                         .unwrap_or("???");
            let val = cap.name("val")
                         .unwrap_or("???");
            if !keys.contains_key(key) {
                keys.insert(key.to_string(), val.to_string());
            }
        }
    }

    keys
}

#[derive(Debug)]
#[derive(Clone)]
struct ValueAndSource {
    value: String,
    source: String,
}

fn insert_if_not_exist(dict: &mut HashMap<String, ValueAndSource>,
                       keys_from_files: &mut HashMap<String, Vec<String>>,
                       matches: HashMap<String, String>,
                       filename: String) {
    for key_val in matches.iter() {
        let key: String = (*key_val.0).clone();
        let val: String = (*key_val.1).clone();
        match dict.entry(key.clone()) {
            Vacant(e) => {
                e.insert(ValueAndSource {
                    value: val,
                    source: filename.clone(),
                });
                match keys_from_files.entry(filename.clone()) {
                    Vacant(k) => {
                        k.insert(vec![key.clone()]);
                    }
                    Occupied(mut k) => {
                        k.get_mut().push(key.clone());
                    }
                }

                // keys_from_files.insert(filename.clone(), key.clone());
            }
            Occupied(e) => {
                println!("[ERROR]: key ({}) in file ({}) tried to override the \
                          same key from file ({})",
                         key,
                         filename,
                         e.get().source);
                process::exit(1);
            }
        }
    }
}

fn read_keyfiles_to_dict(key_filenames: Vec<String>) -> Option<KeysAndSources> {
    let mut dict: HashMap<String, ValueAndSource> = HashMap::new();
    let mut keys_from_files = HashMap::new();
    for filename in key_filenames {
        fs::File::open(&filename)
            .and_then(|f| Ok(extract_keys(f)))
            .map(|v| {
                insert_if_not_exist(&mut dict,
                                    &mut keys_from_files,
                                    v,
                                    filename);
            });
    }

    match dict.len() {
        0 => None,
        _ => {
            Some(KeysAndSources {
                dictionary: Box::new(dict),
                sources: Box::new(keys_from_files),
            })
        }
    }
}

fn create_file(filename: &String) -> fs::File {
    match fs::File::create(filename) {
        Ok(f) => Some(f),
        Err(e) => {
            println_stderr!("[ERROR]: Could not create file: {}\n{}\n",
                            filename,
                            e);
            APP.print_usage_and_panic();
            None
        }
    }
    .unwrap()
}

fn open_file(filename: &String) -> fs::File {
    match fs::File::open(filename) {
        Ok(f) => Some(f),
        Err(e) => {
            println_stderr!("[ERROR]: Could not open file: {}\n{}\n",
                            filename,
                            e);
            APP.print_usage_and_panic();
            None
        }
    }
    .unwrap()
}

struct Settings<'a> {
    key_filenames: Vec<String>,
    prop_filename: &'a String,
    result_filename: &'a String,
    props_first: bool,
}

struct KeysAndSources {
    dictionary: Box<HashMap<String, ValueAndSource>>,
    sources: Box<HashMap<String, Vec<String>>>,
}

fn print_result(used_keys: &mut HashMap<String, Vec<(String, String)>>, prop_filename: &String){

    println!("======== RESULTS ========");

    let env_vars = &used_keys.remove(ENV);
    if env_vars.is_some() {
        let vars = env_vars.clone().unwrap();
        println!("# Environment variables:");
        for v in vars {
            println!("${{{}}}", v.0);
        }
        println!("");
    }

    for src in used_keys.iter() {
        println!("# {}:", src.0);
        let mut keys_src = src.1.clone();
        keys_src.sort();
        for s in keys_src {
            println!("${{{}}}", s.0);
        }
        println!("");
    }

    let missing_vars = used_keys.get(MISSING);
    if missing_vars.is_some() {
        let vars = missing_vars.unwrap();
        println!("Missing variables found in {}:", prop_filename);
        for v in vars {
            println!("${{{}}}", v.0);
        }
        println!("");
        println_stderr!("ERROR: Must replace all found variables\n\n");
        process::exit(1);
    }

    if env_vars.is_none() && used_keys.len() == 0 && missing_vars.is_none() {
       println!("[No variables to substitute!]\n");
    }

    println!("...DONE!");
}

fn calc_result(settings: Settings) {
    let prop_file = open_file(&settings.prop_filename);
    let mut result_buff = Vec::new();
    let mut used_keys = HashMap::new();
    let keys_and_sources = read_keyfiles_to_dict(settings.key_filenames);

    let processed_buffer = process(prop_file,
                                   &mut used_keys,
                                   &keys_and_sources,
                                   settings.props_first);

    for line in processed_buffer.iter() {
        result_buff.push(line.clone());
    }

    let mut result_file = create_file(settings.result_filename);

    for line in result_buff {
        result_file.write_all(&format!("{}\n", line).as_bytes()).unwrap();
    }
    result_file.flush().unwrap();

    print_result(&mut used_keys, settings.prop_filename);
}

fn get_arg_parser() -> ApplicationOptions {
    ApplicationOptions::new(|| {
        let mut o = OptionsAndFlags::new();
        o.parsing_style(getopts::ParsingStyle::StopAtFirstFree);
        o.opt("k",
              "keys",
              "keyfile(s) with variable substitutions, can occur multiple \
               times",
              "FILE(s)",
              getopts::HasArg::Yes,
              getopts::Occur::Multi)
         .unwrap();
        o.opt("p",
              "props-first",
              "properties takes precedence over environment variables \
               (default: off)",
              "",
              getopts::HasArg::No,
              getopts::Occur::Optional)
         .unwrap();
        o.opt("?",
              "help",
              "print this help menu",
              "",
              getopts::HasArg::No,
              getopts::Occur::Optional)
         .unwrap();
        o.opt("V",
              "version",
              "prints current version number",
              "",
              getopts::HasArg::No,
              getopts::Occur::Optional)
         .unwrap();
        Ok(o)
    })
}

struct PropAndResultFilenames<'a> {
    prop_filename: &'a String,
    result_filename: &'a String,
}

fn get_prop_and_result_filename<'a>(free_args: &'a Vec<String>)
                                    -> PropAndResultFilenames<'a> {
    match free_args.len() {
        1 => {
            // Replace vars in the provided file
            PropAndResultFilenames {
                prop_filename: &free_args[0],
                result_filename: &free_args[0],
            }
        }
        2 => {
            // Replace vars into a separate output file
            PropAndResultFilenames {
                prop_filename: &free_args[0],
                result_filename: &free_args[1],
            }
        }
        _ => {
            println_stderr!("Wrong number of file parameters: {}",
                            free_args.len());
            APP.print_usage();
            process::exit(1);
        }
    }
}

fn get_props_first(matches: &Box<getopts::Matches>) -> bool {
    if matches.opt_present("p") {
        true
    } else {
        false
    }

}

fn get_parsed_args() -> Box<getopts::Matches> {
    Box::new(APP.parse().unwrap_or_else(|e| {
        println_stderr!("{}", e);
        process::exit(1);
    }))
}

// Singletons
lazy_static! {
    static ref APP: ApplicationOptions = get_arg_parser();
}

#[allow(dead_code)]
fn main() {
    let matches = get_parsed_args();

    if matches.opt_present("V") {
        println!("proper: v{}", VERSION);
        process::exit(0);
    }

    let prop_and_result_filenames = get_prop_and_result_filename(&matches.free);
    calc_result(Settings {
        key_filenames: matches.opt_strs("k"),
        prop_filename: prop_and_result_filenames.prop_filename,
        result_filename: prop_and_result_filenames.result_filename,
        props_first: get_props_first(&matches),
    });
}

#[cfg(test)]
mod tests {
    use super::replace_var;
    use std::collections::HashMap;
    use super::ValueAndSource;

    struct TestData {
        key_map: HashMap<String, ValueAndSource>,
        used_keys: HashMap<&'static str, Vec<(String, String)>>,
        result: Vec<String>,
    }

    fn create_test_data(keys: Vec<(&str, &str)>) -> TestData {
        let mut test_data = TestData {
            key_map: HashMap::new(),
            used_keys: HashMap::new(),
            result: Vec::new(),
        };

        for k in keys {
            test_data.key_map.insert(k.0.to_string(),
                                     ValueAndSource {
                                         value: k.1.to_string(),
                                         source: "".to_string(),
                                     });
        }

        test_data
    }

    #[test]
    fn can_sub_single_var() {
        let mut test_data = create_test_data(vec![("val", "world!")]);
        replace_var("Hello ${val}".to_string(),
                    &Some(test_data.key_map),
                    true,
                    &mut test_data.used_keys,
                    &mut test_data.result);
        assert!(test_data.result.len() == 1);
        assert!(test_data.used_keys[super::KEY_FILE][0].0 == "val".to_string());
        assert!(test_data.result[0] == "Hello world!".to_string());
    }
}
<commit_msg>Changed a warning description<commit_after>extern crate regex;
extern crate getopts;
#[macro_use]
extern crate lazy_static;
use std::collections::HashMap;
use std::{env, fs, io};
use std::io::prelude::*;
mod options;
use options::*;
use std::process;
use std::collections::hash_map::Entry::{Occupied, Vacant};

macro_rules! regex { ($re:expr) => { ::regex::Regex::new($re).unwrap() } }
macro_rules! println_stderr(
    ($($arg:tt)*) => (
        match writeln!(&mut ::std::io::stderr(), $($arg)* ) {
            Ok(_) => {},
            Err(x) => panic!("Unable to write to stderr: {}", x),
        }
        )
    );

const VERSION: &'static str = env!("CARGO_PKG_VERSION");
const KEY_FILE: &'static str = "KEY_FILE";
const ENV: &'static str = "ENV";
const MISSING: &'static str = "Missing properties";

fn replace_var(line: String,
               keys: &Option<KeysAndSources>,
               props_first: bool,
               used_keys: &mut HashMap<String, Vec<(String, String)>>,
               result: &mut Vec<String>) {
    let re = regex!(r"(?P<full>\$\{\s*(?P<var>[^\}]*)\s*\})");
    let mut str_line: String = line.clone();

    for cap in re.captures_iter(&*str_line.clone()) {
        cap.name("var")
           .map(|v| {
               let env_prop = env_or_prop(v.trim(), props_first, keys)
                                  .map(|v2| {
                                      str_line = re.replace(&*str_line,
                                                            &*(v2.1.value));
                                      (v2.0, v2.1)
                                  });

               env_prop.and_then(|v2| {
                   used_keys.get_mut(&v2.1.source)
                            .and_then(|r| {
                                Some(r.push((v2.0.clone(), v2.1.value.clone())))
                            })
                            .or_else(|| {
                                used_keys.insert(v2.1.source.clone(),
                                                 vec![(v2.0.clone(),
                                                       v2.1.value.clone())]);
                                Some(())
                            })
               })
           });
    }
    result.push(str_line.to_string());
}

fn no_property(key: &str, env: Option<String>) -> Option<(String, ValueAndSource)> {
    match env {
        Some(e) => {
            Some((key.to_string(),
                  ValueAndSource {
                value: e.to_string(),
                source: ENV.to_string(),
            }))
        }
        None => {
            Some((key.to_string(),
                  ValueAndSource {
                value: key.to_string(),
                source: MISSING.to_string(),
            }))
        }
    }
}

fn env_or_prop(key: &str,
               props_first: bool,
               keys: &Option<KeysAndSources>)
               -> Option<(String, ValueAndSource)> {
    match keys {
        &Some(ref k) => {
            let dict = &(k.dictionary);
            let prop = dict.get(key).and_then(|v| Some(v.clone()));
            let env = env::var_os(key)
                          .map(|v2| v2.clone().into_string().ok().unwrap());
            match (prop, props_first) {
                (Some(prop), true) => {
                    Some((key.to_string(),
                          ValueAndSource {
                        value: prop.value.clone(),
                        source: prop.source.clone(),
                    }))
                }
                (Some(prop), false) => {
                    match env {
                        Some(e) => {
                            println!("[WARNING]: key \"${{{}}}\" from \"{}\" was overridden by an \
                                      environment variable.",
                                     key.to_string(),
                                     prop.source);
                            Some((key.to_string(),
                                  ValueAndSource {
                                value: e.to_string(),
                                source: ENV.to_string(),
                            }))
                        }
                        None => {
                            Some((key.to_string(),
                                  ValueAndSource {
                                value: prop.value.clone(),
                                source: prop.source.clone(),
                            }))
                        }
                    }
                }
                (None, true) => {
                    no_property(key, env)
                }
                (None, false) => {
                    no_property(key, env)
                }
            }
        }
        &None => {
            match env::var_os(key)
                      .map(|v| v.into_string().ok())
                      .unwrap_or(None) {
                Some(e) => {
                    Some((key.to_string(),
                          ValueAndSource {
                        value: e.to_string(),
                        source: ENV.to_string(),
                    }))
                }
                None => {
                    Some((key.to_string(),
                          ValueAndSource {
                        value: key.to_string(),
                        source: MISSING.to_string(),
                    }))
                }
            }
        }
    }
}

fn process(file: fs::File,
           used_keys: &mut HashMap<String, Vec<(String, String)>>,
           dict: &Option<KeysAndSources>,
           props_first: bool)
           -> Vec<String> {
    let buf_reader = io::BufReader::new(&file);
    let mut result: Vec<String> = vec![];

    for line in buf_reader.lines() {
        replace_var(line.unwrap(), &dict, props_first, used_keys, &mut result);
    }

    result
}

fn extract_keys(file: fs::File) -> HashMap<String, String> {
    let key_regex = regex!(r"^\s*(?P<key>\S*)\s*=\s*(?P<val>.*)");
    let buf_reader = io::BufReader::new(&file);
    let mut keys = HashMap::new();

    for line in buf_reader.lines().enumerate() {
        for cap in key_regex.captures_iter(&*line.1
                                                 .ok()
                                                 .expect("Could not read \
                                                          line")) {
            let key = cap.name("key")
                         .unwrap_or("???");
            let val = cap.name("val")
                         .unwrap_or("???");
            if !keys.contains_key(key) {
                keys.insert(key.to_string(), val.to_string());
            }
        }
    }

    keys
}

#[derive(Debug)]
#[derive(Clone)]
struct ValueAndSource {
    value: String,
    source: String,
}

fn insert_if_not_exist(dict: &mut HashMap<String, ValueAndSource>,
                       keys_from_files: &mut HashMap<String, Vec<String>>,
                       matches: HashMap<String, String>,
                       filename: String) {
    for key_val in matches.iter() {
        let key: String = (*key_val.0).clone();
        let val: String = (*key_val.1).clone();
        match dict.entry(key.clone()) {
            Vacant(e) => {
                e.insert(ValueAndSource {
                    value: val,
                    source: filename.clone(),
                });
                match keys_from_files.entry(filename.clone()) {
                    Vacant(k) => {
                        k.insert(vec![key.clone()]);
                    }
                    Occupied(mut k) => {
                        k.get_mut().push(key.clone());
                    }
                }

                // keys_from_files.insert(filename.clone(), key.clone());
            }
            Occupied(e) => {
                println!("[ERROR]: key ({}) in file ({}) tried to override \
                          the same key from file ({})",
                         key,
                         filename,
                         e.get().source);
                process::exit(1);
            }
        }
    }
}

fn read_keyfiles_to_dict(key_filenames: Vec<String>) -> Option<KeysAndSources> {
    let mut dict: HashMap<String, ValueAndSource> = HashMap::new();
    let mut keys_from_files = HashMap::new();
    for filename in key_filenames {
        fs::File::open(&filename)
            .and_then(|f| Ok(extract_keys(f)))
            .map(|v| {
                insert_if_not_exist(&mut dict,
                                    &mut keys_from_files,
                                    v,
                                    filename);
            });
    }

    match dict.len() {
        0 => None,
        _ => {
            Some(KeysAndSources {
                dictionary: Box::new(dict),
                sources: Box::new(keys_from_files),
            })
        }
    }
}

fn create_file(filename: &String) -> fs::File {
    match fs::File::create(filename) {
        Ok(f) => Some(f),
        Err(e) => {
            println_stderr!("[ERROR]: Could not create file: {}\n{}\n",
                            filename,
                            e);
            APP.print_usage_and_panic();
            None
        }
    }
    .unwrap()
}

fn open_file(filename: &String) -> fs::File {
    match fs::File::open(filename) {
        Ok(f) => Some(f),
        Err(e) => {
            println_stderr!("[ERROR]: Could not open file: {}\n{}\n",
                            filename,
                            e);
            APP.print_usage_and_panic();
            None
        }
    }
    .unwrap()
}

struct Settings<'a> {
    key_filenames: Vec<String>,
    prop_filename: &'a String,
    result_filename: &'a String,
    props_first: bool,
}

struct KeysAndSources {
    dictionary: Box<HashMap<String, ValueAndSource>>,
    sources: Box<HashMap<String, Vec<String>>>,
}

fn print_result(used_keys: &mut HashMap<String, Vec<(String, String)>>,
                prop_filename: &String) {

    println!("======== RESULTS ========");

    let env_vars = &used_keys.remove(ENV);
    if env_vars.is_some() {
        let vars = env_vars.clone().unwrap();
        println!("# Environment variables:");
        for v in vars {
            println!("${{{}}}", v.0);
        }
        println!("");
    }

    for src in used_keys.iter() {
        println!("# {}:", src.0);
        let mut keys_src = src.1.clone();
        keys_src.sort();
        for s in keys_src {
            println!("${{{}}}", s.0);
        }
        println!("");
    }

    let missing_vars = used_keys.get(MISSING);
    if missing_vars.is_some() {
        let vars = missing_vars.unwrap();
        println!("Missing variables found in {}:", prop_filename);
        for v in vars {
            println!("${{{}}}", v.0);
        }
        println!("");
        println_stderr!("ERROR: Must replace all found variables\n\n");
        process::exit(1);
    }

    if env_vars.is_none() && used_keys.len() == 0 && missing_vars.is_none() {
        println!("[No variables to substitute!]\n");
    }

    println!("...DONE!");
}

fn calc_result(settings: Settings) {
    let prop_file = open_file(&settings.prop_filename);
    let mut result_buff = Vec::new();
    let mut used_keys = HashMap::new();
    let keys_and_sources = read_keyfiles_to_dict(settings.key_filenames);

    let processed_buffer = process(prop_file,
                                   &mut used_keys,
                                   &keys_and_sources,
                                   settings.props_first);

    for line in processed_buffer.iter() {
        result_buff.push(line.clone());
    }

    let mut result_file = create_file(settings.result_filename);

    for line in result_buff {
        result_file.write_all(&format!("{}\n", line).as_bytes()).unwrap();
    }
    result_file.flush().unwrap();

    print_result(&mut used_keys, settings.prop_filename);
}

fn get_arg_parser() -> ApplicationOptions {
    ApplicationOptions::new(|| {
        let mut o = OptionsAndFlags::new();
        o.parsing_style(getopts::ParsingStyle::StopAtFirstFree);
        o.opt("k",
              "keys",
              "keyfile(s) with variable substitutions, can occur multiple \
               times",
              "FILE(s)",
              getopts::HasArg::Yes,
              getopts::Occur::Multi)
         .unwrap();
        o.opt("p",
              "props-first",
              "properties takes precedence over environment variables \
               (default: off)",
              "",
              getopts::HasArg::No,
              getopts::Occur::Optional)
         .unwrap();
        o.opt("?",
              "help",
              "print this help menu",
              "",
              getopts::HasArg::No,
              getopts::Occur::Optional)
         .unwrap();
        o.opt("V",
              "version",
              "prints current version number",
              "",
              getopts::HasArg::No,
              getopts::Occur::Optional)
         .unwrap();
        Ok(o)
    })
}

struct PropAndResultFilenames<'a> {
    prop_filename: &'a String,
    result_filename: &'a String,
}

fn get_prop_and_result_filename<'a>(free_args: &'a Vec<String>)
                                    -> PropAndResultFilenames<'a> {
    match free_args.len() {
        1 => {
            // Replace vars in the provided file
            PropAndResultFilenames {
                prop_filename: &free_args[0],
                result_filename: &free_args[0],
            }
        }
        2 => {
            // Replace vars into a separate output file
            PropAndResultFilenames {
                prop_filename: &free_args[0],
                result_filename: &free_args[1],
            }
        }
        _ => {
            println_stderr!("Wrong number of file parameters: {}",
                            free_args.len());
            APP.print_usage();
            process::exit(1);
        }
    }
}

fn get_props_first(matches: &Box<getopts::Matches>) -> bool {
    if matches.opt_present("p") {
        true
    } else {
        false
    }

}

fn get_parsed_args() -> Box<getopts::Matches> {
    Box::new(APP.parse().unwrap_or_else(|e| {
        println_stderr!("{}", e);
        process::exit(1);
    }))
}

// Singletons
lazy_static! {
    static ref APP: ApplicationOptions = get_arg_parser();
}

#[allow(dead_code)]
fn main() {
    let matches = get_parsed_args();

    if matches.opt_present("V") {
        println!("proper: v{}", VERSION);
        process::exit(0);
    }

    let prop_and_result_filenames = get_prop_and_result_filename(&matches.free);
    calc_result(Settings {
        key_filenames: matches.opt_strs("k"),
        prop_filename: prop_and_result_filenames.prop_filename,
        result_filename: prop_and_result_filenames.result_filename,
        props_first: get_props_first(&matches),
    });
}

#[cfg(test)]
mod tests {
    use super::replace_var;
    use std::collections::HashMap;
    use super::ValueAndSource;

    struct TestData {
        key_map: HashMap<String, ValueAndSource>,
        used_keys: HashMap<&'static str, Vec<(String, String)>>,
        result: Vec<String>,
    }

    fn create_test_data(keys: Vec<(&str, &str)>) -> TestData {
        let mut test_data = TestData {
            key_map: HashMap::new(),
            used_keys: HashMap::new(),
            result: Vec::new(),
        };

        for k in keys {
            test_data.key_map.insert(k.0.to_string(),
                                     ValueAndSource {
                                         value: k.1.to_string(),
                                         source: "".to_string(),
                                     });
        }

        test_data
    }

    #[test]
    fn can_sub_single_var() {
        let mut test_data = create_test_data(vec![("val", "world!")]);
        replace_var("Hello ${val}".to_string(),
                    &Some(test_data.key_map),
                    true,
                    &mut test_data.used_keys,
                    &mut test_data.result);
        assert!(test_data.result.len() == 1);
        assert!(test_data.used_keys[super::KEY_FILE][0].0 == "val".to_string());
        assert!(test_data.result[0] == "Hello world!".to_string());
    }
}
<|endoftext|>


<commit_before>#![feature(conservative_impl_trait)]

extern crate curl;
extern crate env_logger;
extern crate futures;
extern crate tokio_core;
extern crate tokio_curl;
extern crate threadpool;
#[macro_use]
extern crate clap;
extern crate chrono;

mod args;

use args::{Config, parse_args};

use curl::easy::Easy;
use futures::{Async, Future, lazy, Poll, stream};
use futures::future::{ok, join_all, FutureResult, loop_fn, Loop};
use futures::sync::mpsc::channel;
use futures::{Stream, Sink};
use tokio_core::reactor::{Core, Handle};
use tokio_curl::{Perform, PerformError, Session};
use std::rc::Rc;
use std::cell::RefCell;
use threadpool::ThreadPool;
use std::sync::{Arc, Barrier};
use futures::stream::FuturesUnordered;
use chrono::*;

struct Worker {
    num_requests: u32,
}

impl Worker {
    pub fn new() -> Self {
        Worker { num_requests: 0}
    }

    fn send_request(mut self, easy_request: Easy, session: &Session) -> impl Future<Item=(Self, Easy), Error=PerformError> {
        session.perform(easy_request)
            .and_then(move |easy| {
                self.num_requests += 1;
                ok((self, easy))
            })
    }
}

/// Delegates and manages all the work.
struct Boss {
    /// Number of requests successfully completed
    pub requests_completed: usize,
    /// Number of active connections
    pub connections: usize,

    thread_pool: ThreadPool,
    num_threads: usize,
}

impl Boss {
    pub fn new(num_threads: usize) -> Self {
        Boss {
            requests_completed: 0,
            connections: 0,
            thread_pool: ThreadPool::new(num_threads),
            num_threads: num_threads,
        }
    }

    /// Should return a future
    pub fn start_workforce(&self, desired_connections: usize, url: String, duration: Duration) {
        // create a barrier that wait all jobs plus the starter thread
        let barrier = Arc::new(Barrier::new(self.num_threads + 1));
        for _ in 0..self.num_threads {
            let barrier = barrier.clone();
            // let tx = tx.clone();
            let url = url.clone();
            self.thread_pool.execute(move || {

                let mut lp = Core::new().unwrap();
                let start_time = Local::now();
                let wanted_end_time = start_time + duration;
                let session = Session::new(lp.handle());

                let iterator = (0..desired_connections).map(|_| {
                    let mut easy_request = Easy::new();
                    easy_request.get(true).unwrap();
                    easy_request.url(&url).unwrap();
                    loop_fn((Worker::new(), easy_request), |(worker, easy)| {
                        worker.send_request(easy, &session)
                            .and_then(|state| {
                                let now_time = Local::now();
                                if now_time < wanted_end_time {
                                    Ok(Loop::Continue(state))
                                } else {
                                    Ok(Loop::Break(state))
                                }
                            })
                    })
                });
                let future = stream::futures_unordered(iterator)
                    .fold(0, |acc, (res, _)| ok(acc + res.num_requests));
                let res = lp.run(future).unwrap();
                println!("{:?}", res);

                // then wait for the other threads
                barrier.wait();
            });
        }
        // wait for the threads to finish the work
        barrier.wait();
    }
}

fn start(config: Config) {
    let boss = Boss::new(config.num_threads);
    boss.start_workforce(config.num_connections, config.url, config.duration)
}

fn main() {
    env_logger::init().unwrap();
    let config = parse_args();
    start(config);
}
<commit_msg>cleanup<commit_after>#![feature(conservative_impl_trait)]

extern crate curl;
extern crate env_logger;
extern crate futures;
extern crate tokio_core;
extern crate tokio_curl;
extern crate threadpool;
#[macro_use]
extern crate clap;
extern crate chrono;

mod args;

use args::{Config, parse_args};

use curl::easy::Easy;
use futures::{Future, stream};
use futures::future::{ok, loop_fn, Loop};
use futures::{Stream};
use tokio_core::reactor::{Core};
use tokio_curl::{PerformError, Session};
use threadpool::ThreadPool;
use std::sync::{Arc, Barrier};
use chrono::*;

struct Worker {
    num_requests: u32,
}

impl Worker {
    pub fn new() -> Self {
        Worker { num_requests: 0}
    }

    fn send_request(mut self, easy_request: Easy, session: &Session) -> impl Future<Item=(Self, Easy), Error=PerformError> {
        session.perform(easy_request)
            .and_then(move |easy| {
                self.num_requests += 1;
                ok((self, easy))
            })
    }
}

/// Delegates and manages all the work.
struct Boss {
    /// Number of requests successfully completed
    pub requests_completed: usize,
    /// Number of active connections
    pub connections: usize,

    thread_pool: ThreadPool,
    num_threads: usize,
}

impl Boss {
    pub fn new(num_threads: usize) -> Self {
        Boss {
            requests_completed: 0,
            connections: 0,
            thread_pool: ThreadPool::new(num_threads),
            num_threads: num_threads,
        }
    }

    pub fn start_workforce(&self, desired_connections: usize, url: String, duration: Duration) {
        // create a barrier that wait all jobs plus the starter thread
        let barrier = Arc::new(Barrier::new(self.num_threads + 1));
        for _ in 0..self.num_threads {
            let barrier = barrier.clone();
            let url = url.clone();
            self.thread_pool.execute(move || {

                let mut lp = Core::new().unwrap();
                let start_time = Local::now();
                let wanted_end_time = start_time + duration;
                let session = Session::new(lp.handle());

                let iterator = (0..desired_connections).map(|_| {
                    let mut easy_request = Easy::new();
                    easy_request.get(true).unwrap();
                    easy_request.url(&url).unwrap();
                    loop_fn((Worker::new(), easy_request), |(worker, easy)| {
                        worker.send_request(easy, &session)
                            .and_then(|state| {
                                let now_time = Local::now();
                                if now_time < wanted_end_time {
                                    Ok(Loop::Continue(state))
                                } else {
                                    Ok(Loop::Break(state))
                                }
                            })
                    })
                });
                let future = stream::futures_unordered(iterator)
                    .fold(0, |acc, (res, _)| ok(acc + res.num_requests));
                let res = lp.run(future).unwrap();
                println!("{:?}", res);

                // then wait for the other threads
                barrier.wait();
            });
        }
        // wait for the threads to finish the work
        barrier.wait();
    }
}

fn start(config: Config) {
    let boss = Boss::new(config.num_threads);
    boss.start_workforce(config.num_connections, config.url, config.duration)
}

fn main() {
    env_logger::init().unwrap();
    let config = parse_args();
    start(config);
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 Nikita Pekin and the smexybot contributors
// See the README.md file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![cfg_attr(feature = "nightly", feature(proc_macro))]
#![cfg_attr(feature = "clippy", feature(plugin))]
#![cfg_attr(feature = "clippy", plugin(clippy))]
#![warn(missing_copy_implementations,
        missing_debug_implementations,
        missing_docs,
        trivial_casts,
        trivial_numeric_casts,
        unused_extern_crates,
        unused_import_braces,
        unused_qualifications)]
#![deny(missing_docs, non_camel_case_types, unsafe_code)]
#![cfg_attr(not(feature = "nightly"), deny(warnings))]
#![cfg_attr(feature="clippy", warn(
        cast_possible_truncation,
        cast_possible_wrap,
        cast_precision_loss,
        cast_sign_loss,
        mut_mut,
        wrong_pub_self_convention))]
// This allows us to use `unwrap` on `Option` values when compiling in test mode
// (because using it in tests is idiomatic).
#![cfg_attr(all(not(test), feature="clippy"), warn(result_unwrap_used))]

//! Smexybot is a general-purpose [Discord](https://discordapp.com/) bot written
//! in [Rust](https://www.rust-lang.org/). It is built upon the
//! [serenity.rs](https://github.com/zeyla/serenity.rs) Discord API.

extern crate chrono;
extern crate env_logger;
extern crate hyper;
#[cfg(any(feature = "roll", feature = "wolfram", feature = "xkcd"))]
#[macro_use]
extern crate lazy_static;
#[macro_use]
extern crate log;
extern crate rand;
extern crate serde;
#[cfg(feature = "nightly")]
#[macro_use]
extern crate serde_derive;
extern crate serde_json;
extern crate serenity;
extern crate url;

mod command;
mod config;
mod counter;
mod error;
mod util;

use chrono::{DateTime, UTC};
use config::Config;
use counter::CommandCounter;
use serenity::Client;
use serenity::client::LoginType;
use serenity::ext::framework::Framework;
use std::collections::HashMap;
use std::env;
use util::{check_msg, timestamp_to_string};

const RATE_LIMIT_MESSAGE: &'static str = "Try this again in %time% seconds.";

lazy_static! {
    static ref CONFIG: Config = Config::new(Some("config.json"));
    static ref UPTIME: DateTime<UTC> = UTC::now();
}

fn main() {
    // Initialize the `env_logger` to provide logging output.
    env_logger::init().expect("Failed to initialize env_logger");

    // Initialize the `UPTIME` variable.
    debug!("Initialized at: {}", timestamp_to_string(&*UPTIME));

    // Create a client for a user.
    let (_, mut client) = login();

    {
        let mut data = client.data.lock().expect("Failed to lock client data");
        data.insert::<CommandCounter>(HashMap::default());
    }

    client.on_ready(|_context, ready| {
        let shard_info = if let Some(s) = ready.shard {
            Some(format!("shard {}/{} ", s[0] + 1, s[1]))
        } else {
            None
        };
        println!(
            "Started {}as {}#{}, serving {} guilds",
            shard_info.unwrap_or_else(|| "".to_owned()),
            ready.user.name,
            ready.user.discriminator,
            ready.guilds.len(),
        );
    });

    client.with_framework(build_framework);

    if let Err(err) = client.start_autosharded() {
        error!("Client error: {:?}", err);
    }
}

// Configures the `Framework` used by serenity, and registers the handlers for
// any enabled commands.
fn build_framework(framework: Framework) -> Framework {
    let mut framework = framework.configure(|c| {
            c.rate_limit_message(RATE_LIMIT_MESSAGE)
                .prefix(&CONFIG.command_prefix)
        })
        .before(|context, message, command_name| {
            info!(
            "Got command '{}' from user '{}'",
            command_name,
            message.author.name,
        );

            // Increment the number of times this command has been run. If the
            // command's name does not exist in the counter, add a default value of
            // 0.
            let mut data = context.data.lock().expect("Failed to lock context data");
            let counter = data.get_mut::<CommandCounter>().unwrap();
            let entry = counter.entry(command_name.clone()).or_insert(0);
            *entry += 1;

            true
        })
        .after(|context, _message, command_name, error| {
            if let Err(err) = error {
                check_msg(context.say(&err));
            } else {
                debug!("Processed command '{}'", command_name);
            }
        });

    #[cfg(feature = "fuyu")]
    {
        framework = framework.on("fuyu", command::fuyu::handler);
    }
    #[cfg(feature = "help")]
    {
        use serenity::ext::framework::help_commands;
        framework = framework.command("help", |c| c.exec_help(help_commands::plain));
    }
    #[cfg(feature = "ping")]
    {
        framework = framework.on("ping", command::ping::handler);
    }
    #[cfg(feature = "roll")]
    {
        framework = framework.on("roll", command::roll::handler);
    }
    #[cfg(feature = "stats")]
    {
        framework = framework.on("stats", command::stats::handler);
    }
    #[cfg(feature = "tag")]
    {
        framework = framework.on("tag", command::tag::handler);
    }
    #[cfg(feature = "wolfram")]
    {
        framework = framework.on("wolfram", command::wolfram_alpha::handler);
    }
    #[cfg(feature = "xkcd")]
    {
        framework = framework.on("xkcd", command::xkcd::handler);
    }

    framework
}

// Creates a `Client`.
fn login() -> (LoginType, Client) {
    debug!("Attempting to login");

    if let Ok(bot_token) = env::var("DISCORD_BOT_TOKEN") {
        debug!("Performing bot token login");
        return (LoginType::Bot, Client::login_bot(&bot_token));
    }
    debug!("Skipping bot token login");

    if let Ok(user_token) = env::var("DISCORD_USER_TOKEN") {
        debug!("Performing user token login");
        return (LoginType::User, Client::login_user(&user_token));
    }
    debug!("Skipping user token login");

    panic!("No suitable authentication method found");
}
<commit_msg>chore: fix whitespace<commit_after>// Copyright (c) 2016 Nikita Pekin and the smexybot contributors
// See the README.md file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![cfg_attr(feature = "nightly", feature(proc_macro))]
#![cfg_attr(feature = "clippy", feature(plugin))]
#![cfg_attr(feature = "clippy", plugin(clippy))]
#![warn(missing_copy_implementations,
        missing_debug_implementations,
        missing_docs,
        trivial_casts,
        trivial_numeric_casts,
        unused_extern_crates,
        unused_import_braces,
        unused_qualifications)]
#![deny(missing_docs, non_camel_case_types, unsafe_code)]
#![cfg_attr(not(feature = "nightly"), deny(warnings))]
#![cfg_attr(feature="clippy", warn(
        cast_possible_truncation,
        cast_possible_wrap,
        cast_precision_loss,
        cast_sign_loss,
        mut_mut,
        wrong_pub_self_convention))]
// This allows us to use `unwrap` on `Option` values when compiling in test mode
// (because using it in tests is idiomatic).
#![cfg_attr(all(not(test), feature="clippy"), warn(result_unwrap_used))]

//! Smexybot is a general-purpose [Discord](https://discordapp.com/) bot written
//! in [Rust](https://www.rust-lang.org/). It is built upon the
//! [serenity.rs](https://github.com/zeyla/serenity.rs) Discord API.

extern crate chrono;
extern crate env_logger;
extern crate hyper;
#[cfg(any(feature = "roll", feature = "wolfram", feature = "xkcd"))]
#[macro_use]
extern crate lazy_static;
#[macro_use]
extern crate log;
extern crate rand;
extern crate serde;
#[cfg(feature = "nightly")]
#[macro_use]
extern crate serde_derive;
extern crate serde_json;
extern crate serenity;
extern crate url;

mod command;
mod config;
mod counter;
mod error;
mod util;

use chrono::{DateTime, UTC};
use config::Config;
use counter::CommandCounter;
use serenity::Client;
use serenity::client::LoginType;
use serenity::ext::framework::Framework;
use std::collections::HashMap;
use std::env;
use util::{check_msg, timestamp_to_string};

const RATE_LIMIT_MESSAGE: &'static str = "Try this again in %time% seconds.";

lazy_static! {
    static ref CONFIG: Config = Config::new(Some("config.json"));
    static ref UPTIME: DateTime<UTC> = UTC::now();
}

fn main() {
    // Initialize the `env_logger` to provide logging output.
    env_logger::init().expect("Failed to initialize env_logger");

    // Initialize the `UPTIME` variable.
    debug!("Initialized at: {}", timestamp_to_string(&*UPTIME));

    // Create a client for a user.
    let (_, mut client) = login();

    {
        let mut data = client.data.lock().expect("Failed to lock client data");
        data.insert::<CommandCounter>(HashMap::default());
    }

    client.on_ready(|_context, ready| {
        let shard_info = if let Some(s) = ready.shard {
            Some(format!("shard {}/{} ", s[0] + 1, s[1]))
        } else {
            None
        };
        println!(
            "Started {}as {}#{}, serving {} guilds",
            shard_info.unwrap_or_else(|| "".to_owned()),
            ready.user.name,
            ready.user.discriminator,
            ready.guilds.len(),
        );
    });

    client.with_framework(build_framework);

    if let Err(err) = client.start_autosharded() {
        error!("Client error: {:?}", err);
    }
}

// Configures the `Framework` used by serenity, and registers the handlers for
// any enabled commands.
fn build_framework(framework: Framework) -> Framework {
    let mut framework = framework.configure(|c| {
            c.rate_limit_message(RATE_LIMIT_MESSAGE)
                .prefix(&CONFIG.command_prefix)
        })
        .before(|context, message, command_name| {
            info!(
                "Got command '{}' from user '{}'",
                command_name,
                message.author.name,
            );

            // Increment the number of times this command has been run. If the
            // command's name does not exist in the counter, add a default value of
            // 0.
            let mut data = context.data.lock().expect("Failed to lock context data");
            let counter = data.get_mut::<CommandCounter>().unwrap();
            let entry = counter.entry(command_name.clone()).or_insert(0);
            *entry += 1;

            true
        })
        .after(|context, _message, command_name, error| {
            if let Err(err) = error {
                check_msg(context.say(&err));
            } else {
                debug!("Processed command '{}'", command_name);
            }
        });

    #[cfg(feature = "fuyu")]
    {
        framework = framework.on("fuyu", command::fuyu::handler);
    }
    #[cfg(feature = "help")]
    {
        use serenity::ext::framework::help_commands;
        framework = framework.command("help", |c| c.exec_help(help_commands::plain));
    }
    #[cfg(feature = "ping")]
    {
        framework = framework.on("ping", command::ping::handler);
    }
    #[cfg(feature = "roll")]
    {
        framework = framework.on("roll", command::roll::handler);
    }
    #[cfg(feature = "stats")]
    {
        framework = framework.on("stats", command::stats::handler);
    }
    #[cfg(feature = "tag")]
    {
        framework = framework.on("tag", command::tag::handler);
    }
    #[cfg(feature = "wolfram")]
    {
        framework = framework.on("wolfram", command::wolfram_alpha::handler);
    }
    #[cfg(feature = "xkcd")]
    {
        framework = framework.on("xkcd", command::xkcd::handler);
    }

    framework
}

// Creates a `Client`.
fn login() -> (LoginType, Client) {
    debug!("Attempting to login");

    if let Ok(bot_token) = env::var("DISCORD_BOT_TOKEN") {
        debug!("Performing bot token login");
        return (LoginType::Bot, Client::login_bot(&bot_token));
    }
    debug!("Skipping bot token login");

    if let Ok(user_token) = env::var("DISCORD_USER_TOKEN") {
        debug!("Performing user token login");
        return (LoginType::User, Client::login_user(&user_token));
    }
    debug!("Skipping user token login");

    panic!("No suitable authentication method found");
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate log;
extern crate env_logger;
extern crate irc;
#[macro_use]
extern crate lazy_static;
extern crate regex;
extern crate hyper;
extern crate hubcaps;
extern crate hyper_native_tls;
extern crate serde_json;

use std::fmt;
use std::thread;
use std::collections::HashMap;
use regex::Regex;

use irc::client::prelude::*;

use hyper::Client;
use hyper::net::HttpsConnector;
use hyper_native_tls::NativeTlsClient;
use hubcaps::{Credentials, Github};
use hubcaps::comments::CommentOptions;

fn main() {
    env_logger::init().unwrap();

    let server = IrcServer::new("config.json").unwrap();
    server.identify().unwrap();

    let empty_map = HashMap::new();
    let options = if let Some(ref opts) = server.config().options {
        opts
    } else {
        &empty_map
    };

    // FIXME: Eventually this should support multiple channels, plus
    // options to ask the bot which channels it's in, and which channels
    // it currently has buffers in.  (Then we can do things like ask the
    // bot to reboot itself, but it will only do so if it's not busy.)
    let mut channel_data = ChannelData::new(&options);

    for message in server.iter() {
        let message = message.unwrap(); // panic if there's an error

        match message.command {
            Command::PRIVMSG(ref target, ref msg) => {
                match message.source_nickname() {
                    None => {
                        // FIXME: trailing \n
                        warn!("PRIVMSG without a source! {}", message);
                    }
                    Some(ref source) => {
                        let mynick = server.current_nickname();
                        if target == mynick {
                            handle_bot_command(&server, msg, source, None)
                        } else if target.starts_with('#') {
                            let source_ = String::from(*source);
                            let line = if msg.starts_with("\x01ACTION ") && msg.ends_with("\x01") {
                                ChannelLine {
                                    source: source_,
                                    is_action: true,
                                    message: String::from(&msg[8..msg.len() - 1]),
                                }
                            } else {
                                ChannelLine {
                                    source: source_,
                                    is_action: false,
                                    message: msg.clone(),
                                }
                            };

                            // FIXME: This needs to handle requests in /me
                            match check_command_in_channel(mynick, msg) {
                                Some(ref command) => {
                                    handle_bot_command(&server, command, target, Some(source))
                                }
                                None => {
                                    if let Some(response) = channel_data.add_line(line) {
                                        server.send_privmsg(target, &*response).unwrap();
                                    }
                                }
                            }
                        } else {
                            // FIXME: trailing \n
                            warn!("UNEXPECTED TARGET {} in message {}", target, message);
                        }
                    }
                }
            }
            _ => (),
        }
    }
}

// Take a message in the channel, and see if it was a message sent to
// this bot.
fn check_command_in_channel(mynick: &str, msg: &String) -> Option<String> {
    if !msg.starts_with(mynick) {
        return None;
    }
    let after_nick = &msg[mynick.len()..];
    if !after_nick.starts_with(":") && !after_nick.starts_with(",") {
        return None;
    }
    let after_punct = &after_nick[1..];
    Some(String::from(after_punct.trim_left()))
}

fn handle_bot_command(server: &IrcServer,
                      command: &str,
                      response_target: &str,
                      response_username: Option<&str>) {

    let send_line = |response_username: Option<&str>, line: &str| {
        let adjusted_line = match response_username {
            None => String::from(line),
            Some(username) => String::from(username) + ", " + line,
        };
        server
            .send_privmsg(response_target, &adjusted_line)
            .unwrap();
    };

    if command == "help" {
        send_line(response_username, "The commands I understand are:");
        send_line(None, "  help     Send this message.");
        return;
    }

    send_line(response_username,
              "Sorry, I don't understand that command.  Try 'help'.");
}

struct ChannelLine {
    source: String,
    is_action: bool,
    message: String,
}

struct TopicData {
    topic: String,
    github_url: Option<String>,
    lines: Vec<ChannelLine>,
    resolutions: Vec<String>,
}

struct ChannelData<'opts> {
    current_topic: Option<TopicData>,
    options: &'opts HashMap<String, String>,
}

impl fmt::Display for ChannelLine {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.is_action {
            write!(f, "* {} {}", self.source, self.message)
        } else {
            write!(f, "<{}> {}", self.source, self.message)
        }
    }
}

impl TopicData {
    fn new(topic: &str) -> TopicData {
        let topic_ = String::from(topic);
        TopicData {
            topic: topic_,
            github_url: None,
            lines: vec![],
            resolutions: vec![],
        }
    }
}

impl fmt::Display for TopicData {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.resolutions.len() == 0 {
            try!(write!(f,
                        "The CSS Working Group just discussed {}.\n\n",
                        self.topic));
        } else {
            try!(write!(f,
                        "The CSS Working Group just discussed {}, and agreed to the following resolutions:\n\n",
                        self.topic));
            for resolution in &self.resolutions {
                try!(write!(f, "```\n{}\n```\n\n", resolution));
            }
        }

        try!(write!(f,
                    "The full IRC log of that discussion\n"));
        try!(write!(f, "\n```\n"));
        for line in &self.lines {
            try!(write!(f, "{}\n", line));
        }
        try!(write!(f, "```\n\n"));
        Ok(())
    }
}

fn ci_starts_with(s: &str, prefix: &str) -> bool {
    assert!(prefix.to_lowercase() == prefix);

    s.len() >= prefix.len() && s[0..prefix.len()].to_lowercase() == prefix
}

fn strip_ci_prefix(s: &str, prefix: &str) -> Option<String> {
    if ci_starts_with(s, prefix) {
        Some(String::from(s[prefix.len()..].trim_left()))
    } else {
        None
    }
}

impl<'opts> ChannelData<'opts> {
    fn new(options_: &'opts HashMap<String, String>) -> ChannelData {
        ChannelData {
            current_topic: None,
            options: options_,
        }
    }

    // Returns the response that should be sent to the message over IRC.
    fn add_line(&mut self, line: ChannelLine) -> Option<String> {
        if let Some(ref topic) = strip_ci_prefix(&line.message, "topic:") {
            self.start_topic(topic);
        }
        if line.source == "trackbot" && line.is_action == true &&
           line.message == "is ending a teleconference." {
            self.end_topic();
        }
        match self.current_topic {
            None => None,
            Some(ref mut data) => {
                let new_url_option = extract_github_url(&line.message, self.options);
                let response = match (new_url_option.as_ref(), data.github_url.as_ref()) {
                    (None, _) => None,
                    // FIXME: Add and implement instructions to cancel!
                    (Some(new_url), None) => {
                        Some(format!("OK, I'll post this discussion to {}", new_url))
                    }
                    (Some(new_url), Some(old_url)) if old_url == new_url => None,
                    (Some(new_url), Some(old_url)) => {
                        Some(format!("OK, I'll post this discussion to {} instead of {} like you said before",
                                     new_url,
                                     old_url))
                    }
                };

                if let Some(new_url) = new_url_option {
                    data.github_url = Some(new_url);
                }

                if line.message.starts_with("RESOLUTION") ||
                   line.message.starts_with("RESOLVED") {
                    data.resolutions.push(line.message.clone());
                }

                data.lines.push(line);

                response
            }
        }
    }

    fn start_topic(&mut self, topic: &str) {
        if self.current_topic.is_some() {
            self.end_topic();
        }

        self.current_topic = Some(TopicData::new(topic));
    }

    fn end_topic(&mut self) {
        // TODO: Test the topic boundary code.
        if let Some(topic) = self.current_topic.take() {
            if topic.github_url.is_some() {
                let task = GithubCommentTask::new(topic, self.options);
                task.run();
            }
        }
    }
}

fn extract_github_url(message: &str, options: &HashMap<String, String>) -> Option<String> {
    lazy_static! {
        static ref GITHUB_URL_RE: Regex =
            Regex::new(r"^https://github.com/(?P<repo>[^/]*/[^/]*)/issues/(?P<number>[0-9]+)$")
            .unwrap();
    }
    let ref allowed_repos = options["github_repos_allowed"];
    for prefix in ["topic:", "github topic:"].into_iter() {
        if let Some(ref maybe_url) = strip_ci_prefix(&message, prefix) {
            if let Some(ref caps) = GITHUB_URL_RE.captures(maybe_url) {
                for repo in allowed_repos.split_whitespace() {
                    if caps["repo"] == *repo {
                        return Some(maybe_url.clone());
                    }
                }
            }
        }
    }
    None
}

struct GithubCommentTask {
    data: TopicData,
    github: Github,
}

impl GithubCommentTask {
    fn new(data_: TopicData, options: &HashMap<String, String>) -> GithubCommentTask {
        let github_ =
            Github::new(&*options["github_uastring"],
                        Client::with_connector(HttpsConnector::new(NativeTlsClient::new()
                                                                       .unwrap())),
                        Credentials::Token(options["github_access_token"].clone()));
        GithubCommentTask {
            data: data_,
            github: github_,
        }
    }
    fn run(self) {
        thread::spawn(move || { self.main(); });
    }
    fn main(&self) {
        lazy_static! {
            static ref GITHUB_URL_RE: Regex =
                Regex::new(r"^https://github.com/(?P<owner>[^/]*)/(?P<repo>[^/]*)/issues/(?P<number>[0-9]+)$")
                .unwrap();
        }

        if let Some(ref github_url) = self.data.github_url {
            if let Some(ref caps) = GITHUB_URL_RE.captures(github_url) {
                let repo = self.github
                    .repo(String::from(&caps["owner"]), String::from(&caps["repo"]));
                let issue = repo.issue(caps["number"].parse::<u64>().unwrap());
                let comments = issue.comments();

                let comment_text = format!("{}", self.data);
                comments.create(&CommentOptions { body: comment_text });
            } else {
                warn!("How does {} fail to match now when it matched before?",
                      github_url)
            }
        }
    }
}
<commit_msg>Handle /me commands more thoroughly.<commit_after>#[macro_use]
extern crate log;
extern crate env_logger;
extern crate irc;
#[macro_use]
extern crate lazy_static;
extern crate regex;
extern crate hyper;
extern crate hubcaps;
extern crate hyper_native_tls;
extern crate serde_json;

use std::fmt;
use std::thread;
use std::collections::HashMap;
use regex::Regex;

use irc::client::prelude::*;

use hyper::Client;
use hyper::net::HttpsConnector;
use hyper_native_tls::NativeTlsClient;
use hubcaps::{Credentials, Github};
use hubcaps::comments::CommentOptions;

fn main() {
    env_logger::init().unwrap();

    let server = IrcServer::new("config.json").unwrap();
    server.identify().unwrap();

    let empty_map = HashMap::new();
    let options = if let Some(ref opts) = server.config().options {
        opts
    } else {
        &empty_map
    };

    // FIXME: Eventually this should support multiple channels, plus
    // options to ask the bot which channels it's in, and which channels
    // it currently has buffers in.  (Then we can do things like ask the
    // bot to reboot itself, but it will only do so if it's not busy.)
    let mut channel_data = ChannelData::new(&options);

    for message in server.iter() {
        let message = message.unwrap(); // panic if there's an error

        match message.command {
            Command::PRIVMSG(ref target, ref msg) => {
                match message.source_nickname() {
                    None => {
                        // FIXME: trailing \n
                        warn!("PRIVMSG without a source! {}", message);
                    }
                    Some(ref source) => {
                        let source_ = String::from(*source);
                        let line = if msg.starts_with("\x01ACTION ") && msg.ends_with("\x01") {
                            ChannelLine {
                                source: source_,
                                is_action: true,
                                message: String::from(&msg[8..msg.len() - 1]),
                            }
                        } else {
                            ChannelLine {
                                source: source_,
                                is_action: false,
                                message: msg.clone(),
                            }
                        };
                        let mynick = server.current_nickname();
                        if target == mynick {
                            handle_bot_command(&server, &line.message, source, false, None)
                        } else if target.starts_with('#') {
                            match check_command_in_channel(mynick, &line.message) {
                                Some(ref command) => {
                                    handle_bot_command(&server,
                                                       command,
                                                       target,
                                                       line.is_action,
                                                       Some(source))
                                }
                                None => {
                                    if let Some(response) = channel_data.add_line(line) {
                                        server.send_privmsg(target, &*response).unwrap();
                                    }
                                }
                            }
                        } else {
                            // FIXME: trailing \n
                            warn!("UNEXPECTED TARGET {} in message {}", target, message);
                        }
                    }
                }
            }
            _ => (),
        }
    }
}

// Take a message in the channel, and see if it was a message sent to
// this bot.
fn check_command_in_channel(mynick: &str, msg: &String) -> Option<String> {
    if !msg.starts_with(mynick) {
        return None;
    }
    let after_nick = &msg[mynick.len()..];
    if !after_nick.starts_with(":") && !after_nick.starts_with(",") {
        return None;
    }
    let after_punct = &after_nick[1..];
    Some(String::from(after_punct.trim_left()))
}

fn handle_bot_command(server: &IrcServer,
                      command: &str,
                      response_target: &str,
                      response_is_action: bool,
                      response_username: Option<&str>) {

    let send_line = |response_username: Option<&str>, line: &str| {
        let line_with_nick = match response_username {
            None => String::from(line),
            Some(username) => String::from(username) + ", " + line,
        };
        let adjusted_line = match response_is_action {
            false => line_with_nick,
            true => format!("\x01ACTION {}\x01", line_with_nick),
        };
        server
            .send_privmsg(response_target, &adjusted_line)
            .unwrap();
    };

    if command == "help" {
        send_line(response_username, "The commands I understand are:");
        send_line(None, "  help     Send this message.");
        return;
    }

    send_line(response_username,
              "Sorry, I don't understand that command.  Try 'help'.");
}

struct ChannelLine {
    source: String,
    is_action: bool,
    message: String,
}

struct TopicData {
    topic: String,
    github_url: Option<String>,
    lines: Vec<ChannelLine>,
    resolutions: Vec<String>,
}

struct ChannelData<'opts> {
    current_topic: Option<TopicData>,
    options: &'opts HashMap<String, String>,
}

impl fmt::Display for ChannelLine {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.is_action {
            write!(f, "* {} {}", self.source, self.message)
        } else {
            write!(f, "<{}> {}", self.source, self.message)
        }
    }
}

impl TopicData {
    fn new(topic: &str) -> TopicData {
        let topic_ = String::from(topic);
        TopicData {
            topic: topic_,
            github_url: None,
            lines: vec![],
            resolutions: vec![],
        }
    }
}

impl fmt::Display for TopicData {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.resolutions.len() == 0 {
            try!(write!(f,
                        "The CSS Working Group just discussed {}.\n\n",
                        self.topic));
        } else {
            try!(write!(f,
                        "The CSS Working Group just discussed {}, and agreed to the following resolutions:\n\n",
                        self.topic));
            for resolution in &self.resolutions {
                try!(write!(f, "```\n{}\n```\n\n", resolution));
            }
        }

        try!(write!(f,
                    "The full IRC log of that discussion\n"));
        try!(write!(f, "\n```\n"));
        for line in &self.lines {
            try!(write!(f, "{}\n", line));
        }
        try!(write!(f, "```\n\n"));
        Ok(())
    }
}

fn ci_starts_with(s: &str, prefix: &str) -> bool {
    assert!(prefix.to_lowercase() == prefix);

    s.len() >= prefix.len() && s[0..prefix.len()].to_lowercase() == prefix
}

fn strip_ci_prefix(s: &str, prefix: &str) -> Option<String> {
    if ci_starts_with(s, prefix) {
        Some(String::from(s[prefix.len()..].trim_left()))
    } else {
        None
    }
}

impl<'opts> ChannelData<'opts> {
    fn new(options_: &'opts HashMap<String, String>) -> ChannelData {
        ChannelData {
            current_topic: None,
            options: options_,
        }
    }

    // Returns the response that should be sent to the message over IRC.
    fn add_line(&mut self, line: ChannelLine) -> Option<String> {
        if let Some(ref topic) = strip_ci_prefix(&line.message, "topic:") {
            self.start_topic(topic);
        }
        if line.source == "trackbot" && line.is_action == true &&
           line.message == "is ending a teleconference." {
            self.end_topic();
        }
        match self.current_topic {
            None => None,
            Some(ref mut data) => {
                let new_url_option = extract_github_url(&line.message, self.options);
                let response = match (new_url_option.as_ref(), data.github_url.as_ref()) {
                    (None, _) => None,
                    // FIXME: Add and implement instructions to cancel!
                    (Some(new_url), None) => {
                        Some(format!("OK, I'll post this discussion to {}", new_url))
                    }
                    (Some(new_url), Some(old_url)) if old_url == new_url => None,
                    (Some(new_url), Some(old_url)) => {
                        Some(format!("OK, I'll post this discussion to {} instead of {} like you said before",
                                     new_url,
                                     old_url))
                    }
                };

                if let Some(new_url) = new_url_option {
                    data.github_url = Some(new_url);
                }

                if line.message.starts_with("RESOLUTION") || line.message.starts_with("RESOLVED") {
                    data.resolutions.push(line.message.clone());
                }

                data.lines.push(line);

                response
            }
        }
    }

    fn start_topic(&mut self, topic: &str) {
        if self.current_topic.is_some() {
            self.end_topic();
        }

        self.current_topic = Some(TopicData::new(topic));
    }

    fn end_topic(&mut self) {
        // TODO: Test the topic boundary code.
        if let Some(topic) = self.current_topic.take() {
            if topic.github_url.is_some() {
                let task = GithubCommentTask::new(topic, self.options);
                task.run();
            }
        }
    }
}

fn extract_github_url(message: &str, options: &HashMap<String, String>) -> Option<String> {
    lazy_static! {
        static ref GITHUB_URL_RE: Regex =
            Regex::new(r"^https://github.com/(?P<repo>[^/]*/[^/]*)/issues/(?P<number>[0-9]+)$")
            .unwrap();
    }
    let ref allowed_repos = options["github_repos_allowed"];
    for prefix in ["topic:", "github topic:"].into_iter() {
        if let Some(ref maybe_url) = strip_ci_prefix(&message, prefix) {
            if let Some(ref caps) = GITHUB_URL_RE.captures(maybe_url) {
                for repo in allowed_repos.split_whitespace() {
                    if caps["repo"] == *repo {
                        return Some(maybe_url.clone());
                    }
                }
            }
        }
    }
    None
}

struct GithubCommentTask {
    data: TopicData,
    github: Github,
}

impl GithubCommentTask {
    fn new(data_: TopicData, options: &HashMap<String, String>) -> GithubCommentTask {
        let github_ =
            Github::new(&*options["github_uastring"],
                        Client::with_connector(HttpsConnector::new(NativeTlsClient::new()
                                                                       .unwrap())),
                        Credentials::Token(options["github_access_token"].clone()));
        GithubCommentTask {
            data: data_,
            github: github_,
        }
    }
    fn run(self) {
        thread::spawn(move || { self.main(); });
    }
    fn main(&self) {
        lazy_static! {
            static ref GITHUB_URL_RE: Regex =
                Regex::new(r"^https://github.com/(?P<owner>[^/]*)/(?P<repo>[^/]*)/issues/(?P<number>[0-9]+)$")
                .unwrap();
        }

        if let Some(ref github_url) = self.data.github_url {
            if let Some(ref caps) = GITHUB_URL_RE.captures(github_url) {
                let repo = self.github
                    .repo(String::from(&caps["owner"]), String::from(&caps["repo"]));
                let issue = repo.issue(caps["number"].parse::<u64>().unwrap());
                let comments = issue.comments();

                let comment_text = format!("{}", self.data);
                comments.create(&CommentOptions { body: comment_text });
            } else {
                warn!("How does {} fail to match now when it matched before?",
                      github_url)
            }
        }
    }
}
<|endoftext|>


<commit_before>extern crate mio;
extern crate bytes;
extern crate log4rs;
#[macro_use]
extern crate log;
#[macro_use]
extern crate clap;

use mio::tcp::*;
use mio::EventLoop;
use mio::{TryRead,TryWrite};
use mio::util::Slab;
use std::collections::VecDeque;
use clap::{Arg, App};

#[derive(Debug)]
enum MiChatCommand {
    Broadcast(String),
    NewConnection(TcpStream),
}

struct MiChat {
    connections: Slab<EventHandler>,
}

impl MiChat {
    fn new() -> MiChat {
        MiChat {
            connections: Slab::new(1024),
        }
    }

    fn listen(&mut self, event_loop: &mut mio::EventLoop<MiChat>, listener: TcpListener) {
        let l = EventHandler::Listener(Listener::new(listener));
        let server = self.connections.insert(l).expect("insert listener");
        if let &EventHandler::Listener(ref l) = &self.connections[server] {
            event_loop.register(&l.listener, server).expect("Register listener");
        }
    }

    fn process_action(&mut self, msg: MiChatCommand, event_loop: &mut mio::EventLoop<MiChat>) {
        trace!("{:p}; got {:?}", self, msg);
        match msg {
            MiChatCommand::Broadcast(s) => {
                for c in self.connections.iter_mut() {
                    if let &mut EventHandler::Conn(ref mut c) = c {
                        c.enqueue(&s);
                    }
                }
            },

            MiChatCommand::NewConnection(socket) => {
                let token = self.connections
                    .insert_with(|token| EventHandler::Conn(Connection::new(socket, token)))
                    .expect("token insert");

                if let EventHandler::Conn(ref c) = self.connections[token] {
                    event_loop.register_opt(
                            &c.socket,
                            token,
                            mio::EventSet::readable(),
                            mio::PollOpt::edge() | mio::PollOpt::oneshot())
                        .expect("event loop register");
                }
            }
        }
    }
}

#[derive(Debug)]
struct Connection {
    socket: TcpStream,
    sock_status: mio::EventSet,
    token: mio::Token,
    read_buf: Vec<u8>,
    read_eof: bool,
    failed: bool,
    write_buf: Vec<u8>,
}

impl Connection {
    fn new(socket: TcpStream, token: mio::Token) -> Connection {
        Connection {
            socket: socket,
            sock_status: mio::EventSet::none(),
            token: token,
            read_buf: Vec::with_capacity(1024),
            write_buf: Vec::new(),
            read_eof: false,
            failed: false,
        }
    }
    // Event updates arrive here
    fn handle_event(&mut self, _event_loop: &mut mio::EventLoop<MiChat>, events: mio::EventSet) {
        self.sock_status.insert(events);
        info!("Connection::handle_event: {:?}; this time: {:?}; now: {:?}",
                self.socket.peer_addr(), events, self.sock_status);
    }

    // actions are processed here on down.

    fn process_rules(&mut self, event_loop: &mut mio::EventLoop<MiChat>,
        to_parent: &mut VecDeque<MiChatCommand>) {
        if self.sock_status.is_readable() {
            self.read();
            self.sock_status.remove(mio::EventSet::readable());
        }

        self.process_buffer(to_parent);

        if self.sock_status.is_writable() {
            self.write();
            self.sock_status.remove(mio::EventSet::writable());
        }

        if !self.is_closed() {
            self.reregister(event_loop)
        }
    }

    fn process_buffer(&mut self, to_parent: &mut VecDeque<MiChatCommand>) {
        let mut prev = 0;
        info!("{:?}: Read buffer: {:?}", self.socket.peer_addr(), self.read_buf);
        for n in self.read_buf.iter().enumerate()
                .filter_map(|(i, e)| if *e == '\n' as u8 { Some(i) } else { None } ) {
            info!("{:?}: Pos: {:?}; chunk: {:?}", self.socket.peer_addr(), n, &self.read_buf[prev..n]);
            let s = String::from_utf8_lossy(&self.read_buf[prev..n]).to_string();
            let cmd = MiChatCommand::Broadcast(s);
            info!("Send! {:?}", cmd);
            to_parent.push_back(cmd);
            prev = n+1;
        }
        let remainder = self.read_buf[prev..].to_vec();
        info!("{:?}: read Remainder: {}", self.socket.peer_addr(), remainder.len());
        self.read_buf = remainder;
    }

    fn read(&mut self) {
        let mut abuf = Vec::new();
        match self.socket.try_read_buf(&mut abuf) {
            Ok(Some(0)) => {
                info!("{:?}: EOF!", self.socket.peer_addr());
                self.read_eof = true
            },
            Ok(Some(n)) => {
                info!("{:?}: Read {}bytes", self.socket.peer_addr(), n);
                self.read_buf.extend(abuf);
            },
            Ok(None) => {
                info!("{:?}: Noop!", self.socket.peer_addr());
            },
            Err(e) => {
                error!("got an error trying to read; err={:?}", e);
                self.failed =true;
            }
        }
    }

    fn write(&mut self) {
        match self.socket.try_write(&mut self.write_buf) {
            Ok(Some(n)) => {
                info!("{:?}: Wrote {} of {} in buffer", self.socket.peer_addr(), n,
                    self.write_buf.len());
                self.write_buf = self.write_buf[n..].to_vec();
                info!("{:?}: Now {:?}b", self.socket.peer_addr(), self.write_buf.len());
            },
            Ok(None) => {
                info!("Write unready");
            },
            Err(e) => {
                error!("got an error trying to write; err={:?}", e);
                self.failed = true;
            }
        }
    }

    fn reregister(&mut self, event_loop: &mut mio::EventLoop<MiChat>) {
        let mut flags = mio::EventSet::readable();
        if !self.write_buf.is_empty() {
            flags.insert(mio::EventSet::writable());
        }
        info!("Registering {:?} with {:?}", self, flags);

        event_loop.reregister(
                &self.socket,
                self.token,
                flags,
                mio::PollOpt::oneshot()).expect("EventLoop#reregister")
    }

    fn is_closed(&self) -> bool {
        self.failed || (self.read_eof && self.write_buf.is_empty())
    }

    fn enqueue(&mut self, s: &str) {
        self.write_buf.extend(s.as_bytes());
        self.write_buf.push('\n' as u8);
    }
}

#[derive(Debug)]
struct Listener {
    listener: TcpListener,
    sock_status: mio::EventSet,
}

impl Listener {
    fn new(listener: TcpListener) -> Listener {
        Listener {
            listener: listener,
            sock_status: mio::EventSet::none(),
        }
    }

    fn handle_event(&mut self, _event_loop: &mut mio::EventLoop<MiChat>, events: mio::EventSet) {
        assert!(events.is_readable());
        self.sock_status.insert(events);
        info!("Listener::handle_event: {:?}; this time: {:?}; now: {:?}",
                self.listener.local_addr(), events, self.sock_status);
    }

    fn process_rules(&mut self, event_loop: &mut mio::EventLoop<MiChat>,
            to_parent: &mut VecDeque<MiChatCommand>) {
        info!("the listener socket is ready to accept a connection");
        match self.listener.accept() {
            Ok(Some(socket)) => {
                let cmd = MiChatCommand::NewConnection(socket);
                to_parent.push_back(cmd);

/*
*/

            }
            Ok(None) => {
                info!("the listener socket wasn't actually ready");
            }
            Err(e) => {
                info!("listener.accept() errored: {}", e);
                event_loop.shutdown();
            }
        }
    }

    fn is_closed(&self) -> bool {
        false
    }
}

#[derive(Debug)]
enum EventHandler {
    Listener (Listener),
    Conn (Connection),
}

impl EventHandler {
    fn handle_event(&mut self, _event_loop: &mut mio::EventLoop<MiChat>, events: mio::EventSet) {
        match self {
            &mut EventHandler::Conn(ref mut conn) => conn.handle_event(_event_loop, events),
            &mut EventHandler::Listener(ref mut listener) => listener.handle_event(_event_loop, events)
        }
    }

    fn process_rules(&mut self, event_loop: &mut mio::EventLoop<MiChat>,
        to_parent: &mut VecDeque<MiChatCommand>) {
        match self {
            &mut EventHandler::Conn(ref mut conn) => conn.process_rules(event_loop, to_parent),
            &mut EventHandler::Listener(ref mut listener) => listener.process_rules(event_loop, to_parent)
        }
    }

    fn is_closed(&self) -> bool {
        match self {
            &EventHandler::Conn(ref conn) => conn.is_closed(),
            &EventHandler::Listener(ref listener) => listener.is_closed()
        }
    }

}

impl mio::Handler for MiChat {
    type Timeout = ();
    type Message = ();
    fn ready(&mut self, event_loop: &mut mio::EventLoop<Self>, token: mio::Token, events: mio::EventSet) {
        info!("{:?}: {:?}", token, events);
        match token {
            _ => {
                self.connections[token].handle_event(event_loop, events);
                if self.connections[token].is_closed() {
                    info!("Removing; {:?}", token);
                    self.connections.remove(token);
                }
            }
        }

        let mut parent_actions = VecDeque::new();
        loop {
            for conn in self.connections.iter_mut() {
                conn.process_rules(event_loop, &mut parent_actions);
            }
            // Anything left to process?
            if parent_actions.is_empty() { break; }

            for action in parent_actions.drain(..) {
                self.process_action(action, event_loop);
            }
        }
    }
}

const LOG_FILE: &'static str = "log.toml";

fn main() {
    if let Err(e) = log4rs::init_file(LOG_FILE, Default::default()) {
        panic!("Could not init logger from file {}: {}", LOG_FILE, e);
    }
    let matches = App::new("michat")
        .arg(Arg::with_name("bind").short("l").takes_value(true).required(true))
        .get_matches();

    let address = value_t_or_exit!(matches.value_of("bind"), std::net::SocketAddr);
    let listener = TcpListener::bind(&address).expect("bind");

    let mut event_loop = mio::EventLoop::new().expect("Create event loop");
    let mut service = MiChat::new();

    service.listen(&mut event_loop, listener);
    info!("running michat listener at: {:?}", address);
    event_loop.run(&mut service).expect("Run loop");
}
<commit_msg>Factor out registration.<commit_after>extern crate mio;
extern crate bytes;
extern crate log4rs;
#[macro_use]
extern crate log;
#[macro_use]
extern crate clap;

use mio::tcp::*;
use mio::EventLoop;
use mio::{TryRead,TryWrite};
use mio::util::Slab;
use std::collections::VecDeque;
use clap::{Arg, App};

#[derive(Debug)]
enum MiChatCommand {
    Broadcast(String),
    NewConnection(TcpStream),
}

struct MiChat {
    connections: Slab<EventHandler>,
}

impl MiChat {
    fn new() -> MiChat {
        MiChat {
            connections: Slab::new(1024),
        }
    }

    fn listen(&mut self, event_loop: &mut mio::EventLoop<MiChat>, listener: TcpListener) {
        let l = EventHandler::Listener(Listener::new(listener));
        let token = self.connections.insert(l).expect("insert listener");
        &self.connections[token].register(event_loop, token);
    }

    fn process_action(&mut self, msg: MiChatCommand, event_loop: &mut mio::EventLoop<MiChat>) {
        trace!("{:p}; got {:?}", self, msg);
        match msg {
            MiChatCommand::Broadcast(s) => {
                for c in self.connections.iter_mut() {
                    if let &mut EventHandler::Conn(ref mut c) = c {
                        c.enqueue(&s);
                    }
                }
            },

            MiChatCommand::NewConnection(socket) => {
                let token = self.connections
                    .insert_with(|token| EventHandler::Conn(Connection::new(socket, token)))
                    .expect("token insert");
                &self.connections[token].register(event_loop, token);
            }
        }
    }
}

#[derive(Debug)]
struct Connection {
    socket: TcpStream,
    sock_status: mio::EventSet,
    token: mio::Token,
    read_buf: Vec<u8>,
    read_eof: bool,
    failed: bool,
    write_buf: Vec<u8>,
}

impl Connection {
    fn new(socket: TcpStream, token: mio::Token) -> Connection {
        Connection {
            socket: socket,
            sock_status: mio::EventSet::none(),
            token: token,
            read_buf: Vec::with_capacity(1024),
            write_buf: Vec::new(),
            read_eof: false,
            failed: false,
        }
    }

    fn register(&self, event_loop: &mut mio::EventLoop<MiChat>, token: mio::Token) {
        event_loop.register_opt(
                &self.socket,
                token,
                mio::EventSet::readable(),
                mio::PollOpt::edge() | mio::PollOpt::oneshot())
            .expect("event loop register");
    }

    // Event updates arrive here
    fn handle_event(&mut self, _event_loop: &mut mio::EventLoop<MiChat>, events: mio::EventSet) {
        self.sock_status.insert(events);
        info!("Connection::handle_event: {:?}; this time: {:?}; now: {:?}",
                self.socket.peer_addr(), events, self.sock_status);
    }

    // actions are processed here on down.

    fn process_rules(&mut self, event_loop: &mut mio::EventLoop<MiChat>,
        to_parent: &mut VecDeque<MiChatCommand>) {
        if self.sock_status.is_readable() {
            self.read();
            self.sock_status.remove(mio::EventSet::readable());
        }

        self.process_buffer(to_parent);

        if self.sock_status.is_writable() {
            self.write();
            self.sock_status.remove(mio::EventSet::writable());
        }

        if !self.is_closed() {
            self.reregister(event_loop)
        }
    }

    fn process_buffer(&mut self, to_parent: &mut VecDeque<MiChatCommand>) {
        let mut prev = 0;
        info!("{:?}: Read buffer: {:?}", self.socket.peer_addr(), self.read_buf);
        for n in self.read_buf.iter().enumerate()
                .filter_map(|(i, e)| if *e == '\n' as u8 { Some(i) } else { None } ) {
            info!("{:?}: Pos: {:?}; chunk: {:?}", self.socket.peer_addr(), n, &self.read_buf[prev..n]);
            let s = String::from_utf8_lossy(&self.read_buf[prev..n]).to_string();
            let cmd = MiChatCommand::Broadcast(s);
            info!("Send! {:?}", cmd);
            to_parent.push_back(cmd);
            prev = n+1;
        }
        let remainder = self.read_buf[prev..].to_vec();
        info!("{:?}: read Remainder: {}", self.socket.peer_addr(), remainder.len());
        self.read_buf = remainder;
    }

    fn read(&mut self) {
        let mut abuf = Vec::new();
        match self.socket.try_read_buf(&mut abuf) {
            Ok(Some(0)) => {
                info!("{:?}: EOF!", self.socket.peer_addr());
                self.read_eof = true
            },
            Ok(Some(n)) => {
                info!("{:?}: Read {}bytes", self.socket.peer_addr(), n);
                self.read_buf.extend(abuf);
            },
            Ok(None) => {
                info!("{:?}: Noop!", self.socket.peer_addr());
            },
            Err(e) => {
                error!("got an error trying to read; err={:?}", e);
                self.failed =true;
            }
        }
    }

    fn write(&mut self) {
        match self.socket.try_write(&mut self.write_buf) {
            Ok(Some(n)) => {
                info!("{:?}: Wrote {} of {} in buffer", self.socket.peer_addr(), n,
                    self.write_buf.len());
                self.write_buf = self.write_buf[n..].to_vec();
                info!("{:?}: Now {:?}b", self.socket.peer_addr(), self.write_buf.len());
            },
            Ok(None) => {
                info!("Write unready");
            },
            Err(e) => {
                error!("got an error trying to write; err={:?}", e);
                self.failed = true;
            }
        }
    }

    fn reregister(&mut self, event_loop: &mut mio::EventLoop<MiChat>) {
        let mut flags = mio::EventSet::readable();
        if !self.write_buf.is_empty() {
            flags.insert(mio::EventSet::writable());
        }
        info!("Registering {:?} with {:?}", self, flags);

        event_loop.reregister(
                &self.socket,
                self.token,
                flags,
                mio::PollOpt::oneshot()).expect("EventLoop#reregister")
    }

    fn is_closed(&self) -> bool {
        self.failed || (self.read_eof && self.write_buf.is_empty())
    }

    fn enqueue(&mut self, s: &str) {
        self.write_buf.extend(s.as_bytes());
        self.write_buf.push('\n' as u8);
    }
}

#[derive(Debug)]
struct Listener {
    listener: TcpListener,
    sock_status: mio::EventSet,
}

impl Listener {
    fn new(listener: TcpListener) -> Listener {
        Listener {
            listener: listener,
            sock_status: mio::EventSet::none(),
        }
    }

    fn register(&self, event_loop: &mut mio::EventLoop<MiChat>, token: mio::Token) {
        event_loop.register(&self.listener, token).expect("Register listener");
    }

    fn handle_event(&mut self, _event_loop: &mut mio::EventLoop<MiChat>, events: mio::EventSet) {
        assert!(events.is_readable());
        self.sock_status.insert(events);
        info!("Listener::handle_event: {:?}; this time: {:?}; now: {:?}",
                self.listener.local_addr(), events, self.sock_status);
    }

    fn process_rules(&mut self, event_loop: &mut mio::EventLoop<MiChat>,
            to_parent: &mut VecDeque<MiChatCommand>) {
        info!("the listener socket is ready to accept a connection");
        match self.listener.accept() {
            Ok(Some(socket)) => {
                let cmd = MiChatCommand::NewConnection(socket);
                to_parent.push_back(cmd);

/*
*/

            }
            Ok(None) => {
                info!("the listener socket wasn't actually ready");
            }
            Err(e) => {
                info!("listener.accept() errored: {}", e);
                event_loop.shutdown();
            }
        }
    }

    fn is_closed(&self) -> bool {
        false
    }
}

#[derive(Debug)]
enum EventHandler {
    Listener (Listener),
    Conn (Connection),
}

impl EventHandler {
    fn handle_event(&mut self, _event_loop: &mut mio::EventLoop<MiChat>, events: mio::EventSet) {
        match self {
            &mut EventHandler::Conn(ref mut conn) => conn.handle_event(_event_loop, events),
            &mut EventHandler::Listener(ref mut listener) => listener.handle_event(_event_loop, events)
        }
    }

    fn register(&self, event_loop: &mut mio::EventLoop<MiChat>, token: mio::Token) {
        match self {
            &EventHandler::Conn(ref conn) => conn.register(event_loop, token),
            &EventHandler::Listener(ref listener) => listener.register(event_loop, token)
        }
    }

    fn process_rules(&mut self, event_loop: &mut mio::EventLoop<MiChat>,
        to_parent: &mut VecDeque<MiChatCommand>) {
        match self {
            &mut EventHandler::Conn(ref mut conn) => conn.process_rules(event_loop, to_parent),
            &mut EventHandler::Listener(ref mut listener) => listener.process_rules(event_loop, to_parent)
        }
    }

    fn is_closed(&self) -> bool {
        match self {
            &EventHandler::Conn(ref conn) => conn.is_closed(),
            &EventHandler::Listener(ref listener) => listener.is_closed()
        }
    }
}

impl mio::Handler for MiChat {
    type Timeout = ();
    type Message = ();
    fn ready(&mut self, event_loop: &mut mio::EventLoop<Self>, token: mio::Token, events: mio::EventSet) {
        info!("{:?}: {:?}", token, events);
        match token {
            _ => {
                self.connections[token].handle_event(event_loop, events);
                if self.connections[token].is_closed() {
                    info!("Removing; {:?}", token);
                    self.connections.remove(token);
                }
            }
        }

        let mut parent_actions = VecDeque::new();
        loop {
            for conn in self.connections.iter_mut() {
                conn.process_rules(event_loop, &mut parent_actions);
            }
            // Anything left to process?
            if parent_actions.is_empty() { break; }

            for action in parent_actions.drain(..) {
                self.process_action(action, event_loop);
            }
        }
    }
}

const LOG_FILE: &'static str = "log.toml";

fn main() {
    if let Err(e) = log4rs::init_file(LOG_FILE, Default::default()) {
        panic!("Could not init logger from file {}: {}", LOG_FILE, e);
    }
    let matches = App::new("michat")
        .arg(Arg::with_name("bind").short("l").takes_value(true).required(true))
        .get_matches();

    let address = value_t_or_exit!(matches.value_of("bind"), std::net::SocketAddr);
    let listener = TcpListener::bind(&address).expect("bind");

    let mut event_loop = mio::EventLoop::new().expect("Create event loop");
    let mut service = MiChat::new();

    service.listen(&mut event_loop, listener);
    info!("running michat listener at: {:?}", address);
    event_loop.run(&mut service).expect("Run loop");
}
<|endoftext|>