<commit_before>pub use super::super::compiler::*;

use super::super::filter::comments::CommentsRemover;
use super::super::io::memstream::MemStream;
use super::super::lazy::Lazy;

use regex;
use regex::Regex;

use std::env;
use std::io;
use std::io::{Error, Read};
use std::process::{Command, Stdio};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::mpsc::{Receiver, channel};
use std::thread;

lazy_static! {
	static ref RE_CLANG: regex::bytes::Regex = regex::bytes::Regex::new(
	r"(?i)^(clang(:?\+\+)?)(-\d+\.\d+)?(?:.exe)?$"
	).unwrap();
}

pub struct ClangCompiler {
    toolchains: ToolchainHolder,
}

impl ClangCompiler {
    pub fn new() -> Self {
        ClangCompiler { toolchains: ToolchainHolder::new() }
    }
}

struct ClangToolchain {
    path: PathBuf,
    identifier: Lazy<Option<String>>,
}

impl ClangToolchain {
    pub fn new(path: PathBuf) -> Self {
        ClangToolchain {
            path: path,
            identifier: Lazy::new(),
        }
    }
}

impl Compiler for ClangCompiler {
    fn resolve_toolchain(&self, command: &CommandInfo) -> Option<Arc<Toolchain>> {
        if command.program
            .file_name()
            .map_or(false, |n| RE_CLANG.is_match(n.to_string_lossy().as_bytes())) {
            command.find_executable()
                .and_then(|path| self.toolchains.resolve(&path, |path| Arc::new(ClangToolchain::new(path))))
        } else {
            None
        }
    }

    fn discovery_toolchains(&self) -> Vec<Arc<Toolchain>> {
        env::var_os("PATH")
            .map_or(Vec::new(), |paths| env::split_paths(&paths).collect())
            .iter()
            .filter(|path| path.is_absolute())
            .filter_map(|path| path.read_dir().ok())
            .flat_map(|read_dir| read_dir)
            .filter_map(|entry| entry.ok())
            .filter(|entry| RE_CLANG.is_match(entry.file_name().to_string_lossy().as_bytes()))
            .map(|entry| -> Arc<Toolchain> { Arc::new(ClangToolchain::new(entry.path())) })
            .collect()
    }
}

impl Toolchain for ClangToolchain {
    fn identifier(&self) -> Option<String> {
        self.identifier.get(|| clang_identifier(&self.path))
    }

    fn create_tasks(&self, command: CommandInfo, args: &[String]) -> Result<Vec<CompilationTask>, String> {
        super::prepare::create_tasks(command, args)
    }

    fn preprocess_step(&self, state: &SharedState, task: &CompilationTask) -> Result<PreprocessResult, Error> {
        let mut args = Vec::new();
        args.push("-E".to_string());
        args.push("-x".to_string());
        args.push(task.language.clone());
        args.push("-frewrite-includes".to_string());

        // Make parameters list for preprocessing.
        for arg in task.shared.args.iter() {
            match arg {
                &Arg::Flag { ref scope, ref flag } => {
                    match scope {
                        &Scope::Preprocessor |
                        &Scope::Shared => {
                            args.push("-".to_string() + &flag);
                        }
                        &Scope::Ignore | &Scope::Compiler => {}
                    }
                }
                &Arg::Param { ref scope, ref flag, ref value } => {
                    match scope {
                        &Scope::Preprocessor |
                        &Scope::Shared => {
                            args.push("-".to_string() + &flag);
                            args.push(value.clone());
                        }
                        &Scope::Ignore | &Scope::Compiler => {}
                    }
                }
                &Arg::Input { .. } => {}
                &Arg::Output { .. } => {}
            };
        }

        // Add preprocessor paramters.
        args.push(task.input_source.display().to_string());
        args.push("-o".to_string());
        args.push("-".to_string());

        state.wrap_slow(|| execute(task.shared.command.to_command().args(&args)))
    }

    // Compile preprocessed file.
    fn compile_prepare_step(&self, task: CompilationTask, preprocessed: MemStream) -> Result<CompileStep, Error> {
        let mut args = Vec::new();
        args.push("-x".to_string());
        args.push(task.language.clone());
        for arg in task.shared.args.iter() {
            match arg {
                &Arg::Flag { ref scope, ref flag } => {
                    match scope {
                        &Scope::Compiler | &Scope::Shared => {
                            args.push("-".to_string() + &flag);
                        }
                        &Scope::Ignore |
                        &Scope::Preprocessor => {}
                    }
                }
                &Arg::Param { ref scope, ref flag, ref value } => {
                    match scope {
                        &Scope::Compiler | &Scope::Shared => {
                            args.push("-".to_string() + &flag);
                            args.push(value.clone());
                        }
                        &Scope::Ignore |
                        &Scope::Preprocessor => {}
                    }
                }
                &Arg::Input { .. } => {}
                &Arg::Output { .. } => {}
            };
        }
        Ok(CompileStep::new(task, preprocessed, args, false))
    }

    fn compile_step(&self, state: &SharedState, task: CompileStep) -> Result<OutputInfo, Error> {
        // Run compiler.
        state.wrap_slow(|| {
            Command::new(&self.path)
                .env_clear()
                .arg("-c")
                .args(&task.args)
                .arg("-")
                .arg("-o")
                .arg(task.output_object.as_ref().map_or("-".to_string(), |path| path.display().to_string()))
                .stdin(Stdio::piped())
                .stdout(Stdio::piped())
                .stderr(Stdio::piped())
                .spawn()
                .and_then(|mut child| {
                    try!(task.preprocessed.copy(child.stdin.as_mut().unwrap()));
                    let _ = task.preprocessed;
                    child.wait_with_output()
                })
                .map(|o| OutputInfo::new(o))
        })
    }
}

fn clang_parse_version(base_name: &str, stdout: &str) -> Option<String> {
    lazy_static!{
		static ref RE: Regex = Regex::new(r"^.*clang.*?\((\S+)\).*\nTarget:\s*(\S+)").unwrap();
	}

    RE.captures_iter(&stdout)
        .next()
        .and_then(|cap| {
            Some(format!("{} {} {}",
                         base_name,
                         cap.at(1).unwrap_or(""),
                         cap.at(2).unwrap_or("")))
        })
}

fn clang_identifier(clang: &Path) -> Option<String> {
    clang.file_name()
        .and_then(|file_name| {
            RE_CLANG.captures_iter(file_name.to_string_lossy().as_bytes())
                .next()
                .and_then(|cap| cap.at(1))
                .map(|base_name| String::from_utf8_lossy(base_name).into_owned())
        })
        .and_then(|base_name| {
            Command::new(clang.as_os_str())
                .arg("--version")
                .output()
                .ok()
                .and_then(|output| {
                    match output.status.success() {
                        true => Some(String::from_utf8_lossy(&output.stdout).to_string()),
                        false => None,
                    }
                })
                .and_then(|stdout| clang_parse_version(&base_name, &stdout))
        })
}

fn execute(command: &mut Command) -> Result<PreprocessResult, Error> {
    let mut child = try!(command.stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn());
    drop(child.stdin.take());

    fn read_stdout<T: Read>(stream: Option<T>) -> MemStream {
        stream.map_or(Ok(MemStream::new()), |mut stream| {
                let mut ret = MemStream::new();
                io::copy(&mut stream, &mut ret).map(|_| ret)
            })
            .unwrap_or(MemStream::new())
    }

    fn read_stderr<T: Read + Send + 'static>(stream: Option<T>) -> Receiver<Result<Vec<u8>, Error>> {
        let (tx, rx) = channel();
        match stream {
            Some(mut stream) => {
                thread::spawn(move || {
                    let mut ret = Vec::new();
                    let res = stream.read_to_end(&mut ret).map(|_| ret);
                    tx.send(res).unwrap();
                });
            }
            None => tx.send(Ok(Vec::new())).unwrap(),
        }
        rx
    }

    fn bytes(stream: Receiver<Result<Vec<u8>, Error>>) -> Vec<u8> {
        stream.recv().unwrap().unwrap_or(Vec::new())
    }

    let stdout = read_stdout(child.stdout.take().map(|f| CommentsRemover::new(f)));
    let rx_err = read_stderr(child.stderr.take());
    let status = try!(child.wait());
    let stderr = bytes(rx_err);

    if status.success() {
        Ok(PreprocessResult::Success(stdout))
    } else {
        Ok(PreprocessResult::Failed(OutputInfo {
            status: status.code(),
            stdout: Vec::new(),
            stderr: stderr,
        }))
    }
}
#[cfg(test)]
mod test {
    #[test]
    fn test_ubuntu_14_04_clang_3_5() {
        assert_eq!(super::clang_parse_version("prefix",
r#"Ubuntu clang version 3.5.0-4ubuntu2~trusty2 (tags/RELEASE_350/final) (based on LLVM 3.5.0)
Target: x86_64-pc-linux-gnu
Thread model: posix
"#), Some("prefix tags/RELEASE_350/final x86_64-pc-linux-gnu".to_string()))
    }

    #[test]
    fn test_ubuntu_14_04_clang_3_6() {
        assert_eq!(super::clang_parse_version("prefix",
r#"Ubuntu clang version 3.6.0-2ubuntu1~trusty1 (tags/RELEASE_360/final) (based on LLVM 3.6.0)
Target: x86_64-pc-linux-gnu
Thread model: posix
"#), Some("prefix tags/RELEASE_360/final x86_64-pc-linux-gnu".to_string()))
    }

    #[test]
    fn test_ubuntu_16_04_clang_3_5() {
        assert_eq!(super::clang_parse_version("prefix",
r#"Ubuntu clang version 3.5.2-3ubuntu1 (tags/RELEASE_352/final) (based on LLVM 3.5.2)
Target: x86_64-pc-linux-gnu
Thread model: posix
"#), Some("prefix tags/RELEASE_352/final x86_64-pc-linux-gnu".to_string()))
    }

    #[test]
    fn test_ubuntu_16_04_clang_3_6() {
        assert_eq!(super::clang_parse_version("prefix",
r#"Ubuntu clang version 3.6.2-3ubuntu2 (tags/RELEASE_362/final) (based on LLVM 3.6.2)
Target: x86_64-pc-linux-gnu
Thread model: posix
"#), Some("prefix tags/RELEASE_362/final x86_64-pc-linux-gnu".to_string()))
    }

    #[test]
    fn test_ubuntu_16_04_clang_3_7() {
        assert_eq!(super::clang_parse_version("prefix",
r#"Ubuntu clang version 3.7.1-2ubuntu2 (tags/RELEASE_371/final) (based on LLVM 3.7.1)
Target: x86_64-pc-linux-gnu
Thread model: posix
"#), Some("prefix tags/RELEASE_371/final x86_64-pc-linux-gnu".to_string()))
    }

    #[test]
    fn test_ubuntu_16_04_clang_3_8() {
        assert_eq!(super::clang_parse_version("prefix",
                                              r#"clang version 3.8.0-2ubuntu3 (tags/RELEASE_380/final)
Target: x86_64-pc-linux-gnu
Thread model: posix
InstalledDir: /usr/bin
"#),
                   Some("prefix tags/RELEASE_380/final x86_64-pc-linux-gnu".to_string()))
    }
}
<commit_msg>Fix potential hangup on big clang stderr output<commit_after>pub use super::super::compiler::*;

use super::super::filter::comments::CommentsRemover;
use super::super::io::memstream::MemStream;
use super::super::lazy::Lazy;

use regex;
use regex::Regex;

use std::env;
use std::io;
use std::io::{Error, Read};
use std::process::{Command, Stdio};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::mpsc::{Receiver, channel};
use std::thread;

lazy_static! {
	static ref RE_CLANG: regex::bytes::Regex = regex::bytes::Regex::new(
	r"(?i)^(clang(:?\+\+)?)(-\d+\.\d+)?(?:.exe)?$"
	).unwrap();
}

pub struct ClangCompiler {
    toolchains: ToolchainHolder,
}

impl ClangCompiler {
    pub fn new() -> Self {
        ClangCompiler { toolchains: ToolchainHolder::new() }
    }
}

struct ClangToolchain {
    path: PathBuf,
    identifier: Lazy<Option<String>>,
}

impl ClangToolchain {
    pub fn new(path: PathBuf) -> Self {
        ClangToolchain {
            path: path,
            identifier: Lazy::new(),
        }
    }
}

impl Compiler for ClangCompiler {
    fn resolve_toolchain(&self, command: &CommandInfo) -> Option<Arc<Toolchain>> {
        if command.program
            .file_name()
            .map_or(false, |n| RE_CLANG.is_match(n.to_string_lossy().as_bytes())) {
            command.find_executable()
                .and_then(|path| self.toolchains.resolve(&path, |path| Arc::new(ClangToolchain::new(path))))
        } else {
            None
        }
    }

    fn discovery_toolchains(&self) -> Vec<Arc<Toolchain>> {
        env::var_os("PATH")
            .map_or(Vec::new(), |paths| env::split_paths(&paths).collect())
            .iter()
            .filter(|path| path.is_absolute())
            .filter_map(|path| path.read_dir().ok())
            .flat_map(|read_dir| read_dir)
            .filter_map(|entry| entry.ok())
            .filter(|entry| RE_CLANG.is_match(entry.file_name().to_string_lossy().as_bytes()))
            .map(|entry| -> Arc<Toolchain> { Arc::new(ClangToolchain::new(entry.path())) })
            .collect()
    }
}

impl Toolchain for ClangToolchain {
    fn identifier(&self) -> Option<String> {
        self.identifier.get(|| clang_identifier(&self.path))
    }

    fn create_tasks(&self, command: CommandInfo, args: &[String]) -> Result<Vec<CompilationTask>, String> {
        super::prepare::create_tasks(command, args)
    }

    fn preprocess_step(&self, state: &SharedState, task: &CompilationTask) -> Result<PreprocessResult, Error> {
        let mut args = Vec::new();
        args.push("-E".to_string());
        args.push("-x".to_string());
        args.push(task.language.clone());
        args.push("-frewrite-includes".to_string());

        // Make parameters list for preprocessing.
        for arg in task.shared.args.iter() {
            match arg {
                &Arg::Flag { ref scope, ref flag } => {
                    match scope {
                        &Scope::Preprocessor |
                        &Scope::Shared => {
                            args.push("-".to_string() + &flag);
                        }
                        &Scope::Ignore | &Scope::Compiler => {}
                    }
                }
                &Arg::Param { ref scope, ref flag, ref value } => {
                    match scope {
                        &Scope::Preprocessor |
                        &Scope::Shared => {
                            args.push("-".to_string() + &flag);
                            args.push(value.clone());
                        }
                        &Scope::Ignore | &Scope::Compiler => {}
                    }
                }
                &Arg::Input { .. } => {}
                &Arg::Output { .. } => {}
            };
        }

        // Add preprocessor paramters.
        args.push(task.input_source.display().to_string());
        args.push("-o".to_string());
        args.push("-".to_string());

        state.wrap_slow(|| execute(task.shared.command.to_command().args(&args)))
    }

    // Compile preprocessed file.
    fn compile_prepare_step(&self, task: CompilationTask, preprocessed: MemStream) -> Result<CompileStep, Error> {
        let mut args = Vec::new();
        args.push("-x".to_string());
        args.push(task.language.clone());
        for arg in task.shared.args.iter() {
            match arg {
                &Arg::Flag { ref scope, ref flag } => {
                    match scope {
                        &Scope::Compiler | &Scope::Shared => {
                            args.push("-".to_string() + &flag);
                        }
                        &Scope::Ignore |
                        &Scope::Preprocessor => {}
                    }
                }
                &Arg::Param { ref scope, ref flag, ref value } => {
                    match scope {
                        &Scope::Compiler | &Scope::Shared => {
                            args.push("-".to_string() + &flag);
                            args.push(value.clone());
                        }
                        &Scope::Ignore |
                        &Scope::Preprocessor => {}
                    }
                }
                &Arg::Input { .. } => {}
                &Arg::Output { .. } => {}
            };
        }
        Ok(CompileStep::new(task, preprocessed, args, false))
    }

    fn compile_step(&self, state: &SharedState, task: CompileStep) -> Result<OutputInfo, Error> {
        // Run compiler.
        state.wrap_slow(|| {
            Command::new(&self.path)
                .env_clear()
                .arg("-c")
                .args(&task.args)
                .arg("-")
                .arg("-o")
                .arg(task.output_object.as_ref().map_or("-".to_string(), |path| path.display().to_string()))
                .stdin(Stdio::piped())
                .stdout(Stdio::piped())
                .stderr(Stdio::piped())
                .spawn()
                .and_then(|mut child| {
                    try!(task.preprocessed.copy(child.stdin.as_mut().unwrap()));
                    let _ = task.preprocessed;
                    child.wait_with_output()
                })
                .map(|o| OutputInfo::new(o))
        })
    }
}

fn clang_parse_version(base_name: &str, stdout: &str) -> Option<String> {
    lazy_static!{
		static ref RE: Regex = Regex::new(r"^.*clang.*?\((\S+)\).*\nTarget:\s*(\S+)").unwrap();
	}

    RE.captures_iter(&stdout)
        .next()
        .and_then(|cap| {
            Some(format!("{} {} {}",
                         base_name,
                         cap.at(1).unwrap_or(""),
                         cap.at(2).unwrap_or("")))
        })
}

fn clang_identifier(clang: &Path) -> Option<String> {
    clang.file_name()
        .and_then(|file_name| {
            RE_CLANG.captures_iter(file_name.to_string_lossy().as_bytes())
                .next()
                .and_then(|cap| cap.at(1))
                .map(|base_name| String::from_utf8_lossy(base_name).into_owned())
        })
        .and_then(|base_name| {
            Command::new(clang.as_os_str())
                .arg("--version")
                .output()
                .ok()
                .and_then(|output| {
                    match output.status.success() {
                        true => Some(String::from_utf8_lossy(&output.stdout).to_string()),
                        false => None,
                    }
                })
                .and_then(|stdout| clang_parse_version(&base_name, &stdout))
        })
}

fn execute(command: &mut Command) -> Result<PreprocessResult, Error> {
    let mut child = try!(command.stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn());
    drop(child.stdin.take());

    fn read_stdout<T: Read>(stream: Option<T>) -> MemStream {
        stream.map_or(Ok(MemStream::new()), |mut stream| {
                let mut ret = MemStream::new();
                io::copy(&mut stream, &mut ret).map(|_| ret)
            })
            .unwrap_or(MemStream::new())
    }

    fn read_stderr<T: Read + Send + 'static>(stream: Option<T>) -> Receiver<Result<Vec<u8>, Error>> {
        let (tx, rx) = channel();
        match stream {
            Some(mut stream) => {
                thread::spawn(move || {
                    let mut ret = Vec::new();
                    let res = stream.read_to_end(&mut ret).map(|_| ret);
                    tx.send(res).unwrap();
                });
            }
            None => tx.send(Ok(Vec::new())).unwrap(),
        }
        rx
    }

    fn bytes(stream: Receiver<Result<Vec<u8>, Error>>) -> Vec<u8> {
        stream.recv().unwrap().unwrap_or(Vec::new())
    }

    let rx_err = read_stderr(child.stderr.take());
    let stdout = read_stdout(child.stdout.take().map(|f| CommentsRemover::new(f)));
    let status = try!(child.wait());
    let stderr = bytes(rx_err);

    if status.success() {
        Ok(PreprocessResult::Success(stdout))
    } else {
        Ok(PreprocessResult::Failed(OutputInfo {
            status: status.code(),
            stdout: Vec::new(),
            stderr: stderr,
        }))
    }
}

#[cfg(test)]
mod test {
    #[test]
    fn test_ubuntu_14_04_clang_3_5() {
        assert_eq!(super::clang_parse_version("prefix",
r#"Ubuntu clang version 3.5.0-4ubuntu2~trusty2 (tags/RELEASE_350/final) (based on LLVM 3.5.0)
Target: x86_64-pc-linux-gnu
Thread model: posix
"#), Some("prefix tags/RELEASE_350/final x86_64-pc-linux-gnu".to_string()))
    }

    #[test]
    fn test_ubuntu_14_04_clang_3_6() {
        assert_eq!(super::clang_parse_version("prefix",
r#"Ubuntu clang version 3.6.0-2ubuntu1~trusty1 (tags/RELEASE_360/final) (based on LLVM 3.6.0)
Target: x86_64-pc-linux-gnu
Thread model: posix
"#), Some("prefix tags/RELEASE_360/final x86_64-pc-linux-gnu".to_string()))
    }

    #[test]
    fn test_ubuntu_16_04_clang_3_5() {
        assert_eq!(super::clang_parse_version("prefix",
r#"Ubuntu clang version 3.5.2-3ubuntu1 (tags/RELEASE_352/final) (based on LLVM 3.5.2)
Target: x86_64-pc-linux-gnu
Thread model: posix
"#), Some("prefix tags/RELEASE_352/final x86_64-pc-linux-gnu".to_string()))
    }

    #[test]
    fn test_ubuntu_16_04_clang_3_6() {
        assert_eq!(super::clang_parse_version("prefix",
r#"Ubuntu clang version 3.6.2-3ubuntu2 (tags/RELEASE_362/final) (based on LLVM 3.6.2)
Target: x86_64-pc-linux-gnu
Thread model: posix
"#), Some("prefix tags/RELEASE_362/final x86_64-pc-linux-gnu".to_string()))
    }

    #[test]
    fn test_ubuntu_16_04_clang_3_7() {
        assert_eq!(super::clang_parse_version("prefix",
r#"Ubuntu clang version 3.7.1-2ubuntu2 (tags/RELEASE_371/final) (based on LLVM 3.7.1)
Target: x86_64-pc-linux-gnu
Thread model: posix
"#), Some("prefix tags/RELEASE_371/final x86_64-pc-linux-gnu".to_string()))
    }

    #[test]
    fn test_ubuntu_16_04_clang_3_8() {
        assert_eq!(super::clang_parse_version("prefix",
                                              r#"clang version 3.8.0-2ubuntu3 (tags/RELEASE_380/final)
Target: x86_64-pc-linux-gnu
Thread model: posix
InstalledDir: /usr/bin
"#),
                   Some("prefix tags/RELEASE_380/final x86_64-pc-linux-gnu".to_string()))
    }
}
<|endoftext|>


<commit_before>//! A system for creating and managing IRC client connections.
//!
//! This API provides the ability to create and manage multiple IRC clients that can run on the same
//! thread through the use of a shared event loop. It can also be used to encapsulate the dependency
//! on `tokio` and `futures` in the use of `IrcClient::new_future`. This means that knowledge of
//! those libraries should be unnecessary for the average user. Nevertheless, this API also provides
//! some escape hatches that let advanced users take further advantage of these dependencies.
//!
//! # Example
//! ```no_run
//! # extern crate irc;
//! # use std::default::Default;
//! use irc::client::prelude::*;
//! use irc::error;
//!
//! fn main() {
//!   let config = Config::default();
//!   let mut reactor = IrcReactor::new().unwrap();
//!   let client = reactor.prepare_client_and_connect(&config).unwrap();
//!   reactor.register_client_with_handler(client, process_msg);
//!   reactor.run().unwrap();
//! }
//! # fn process_msg(client: &IrcClient, message: Message) -> error::Result<()> { Ok(()) }
//! ```

use futures::{Future, IntoFuture, Stream};
use futures::future;
use tokio_core::reactor::{Core, Handle};

use client::data::Config;
use client::{IrcClient, IrcClientFuture, PackedIrcClient, Client};
use error;
use proto::Message;

/// A thin wrapper over an event loop.
///
/// An IRC reactor is used to create new IRC clients and to drive the management of all connected
/// clients as the application runs. It can be used to run multiple clients on the same thread, as
/// well as to get better control over error management in an IRC client.
///
/// For a full example usage, see [`irc::client::reactor`](./index.html).
pub struct IrcReactor {
    inner: Core,
    handlers: Vec<Box<Future<Item = (), Error = error::IrcError>>>,
}

impl IrcReactor {
    /// Creates a new reactor.
    pub fn new() -> error::Result<IrcReactor> {
        Ok(IrcReactor {
            inner: Core::new()?,
            handlers: Vec::new(),
        })
    }

    /// Creates a representation of an IRC client that has not yet attempted to connect. In
    /// particular, this representation is as a `Future` that when run will produce a connected
    /// [`IrcClient`](../struct.IrcClient.html).
    ///
    /// # Example
    /// ```no_run
    /// # extern crate irc;
    /// # use std::default::Default;
    /// # use irc::client::prelude::*;
    /// # fn main() {
    /// # let config = Config::default();
    /// let future_client = IrcReactor::new().and_then(|mut reactor| {
    ///     reactor.prepare_client(&config)
    /// });
    /// # }
    /// ```
    pub fn prepare_client<'a>(&mut self, config: &'a Config) -> error::Result<IrcClientFuture<'a>> {
        IrcClient::new_future(self.inner_handle(), config)
    }

    /// Runs an [`IrcClientFuture`](../struct.IrcClientFuture.html), such as one from
    /// `prepare_client` to completion, yielding an [`IrcClient`](../struct.IrcClient.html).
    ///
    /// # Example
    /// ```no_run
    /// # extern crate irc;
    /// # use std::default::Default;
    /// # use irc::client::prelude::*;
    /// # fn main() {
    /// # let config = Config::default();
    /// let client = IrcReactor::new().and_then(|mut reactor| {
    ///     reactor.prepare_client(&config).and_then(|future| {
    ///         reactor.connect_client(future)
    ///     })
    /// });
    /// # }
    /// ```
    pub fn connect_client(&mut self, future: IrcClientFuture) -> error::Result<IrcClient> {
        self.inner.run(future).map(|PackedIrcClient(client, future)| {
            self.register_future(future);
            client
        })
    }

    /// Creates a new [`IrcClient`](../struct.IrcClient.html) from the specified configuration,
    /// connecting immediately. This is guaranteed to be the composition of `prepare_client` and
    /// `connect_client`.
    ///
    /// # Example
    /// ```no_run
    /// # extern crate irc;
    /// # use std::default::Default;
    /// # use irc::client::prelude::*;
    /// # fn main() {
    /// # let config = Config::default();
    /// let client = IrcReactor::new().and_then(|mut reactor| {
    ///     reactor.prepare_client_and_connect(&config)
    /// });
    /// # }
    /// ```
    pub fn prepare_client_and_connect(&mut self, config: &Config) -> error::Result<IrcClient> {
        self.prepare_client(config).and_then(|future| self.connect_client(future))
    }

    /// Registers the given client with the specified message handler. The reactor will store this
    /// setup until the next call to run, where it will be used to process new messages over the
    /// connection indefinitely (or until failure). As registration is consumed by `run`, subsequent
    /// calls to run will require new registration.
    ///
    /// # Example
    /// ```no_run
    /// # extern crate irc;
    /// # use std::default::Default;
    /// # use irc::client::prelude::*;
    /// # fn main() {
    /// # let config = Config::default();
    /// let mut reactor = IrcReactor::new().unwrap();
    /// let client = reactor.prepare_client_and_connect(&config).unwrap();
    /// reactor.register_client_with_handler(client, |client, msg| {
    ///   // Message processing happens here.
    ///   Ok(())
    /// })
    /// # }
    /// ```
    pub fn register_client_with_handler<F, U>(
        &mut self, client: IrcClient, mut handler: F
    ) where F: FnMut(&IrcClient, Message) -> U + 'static,
            U: IntoFuture<Item = (), Error = error::IrcError> + 'static {
        self.handlers.push(Box::new(client.stream().for_each(move |message| {
            handler(&client, message)
        })));
    }

    /// Registers an arbitrary future with this reactor. This is a sort of escape hatch that allows
    /// you to take more control over what runs on the reactor without requiring you to bring in
    /// additional knowledge about `tokio`. It is suspected that `register_client_with_handler` will
    /// be sufficient for most use cases.
    pub fn register_future<F>(
        &mut self, future: F
    ) where F: IntoFuture<Item = (), Error = error::IrcError> + 'static {
        self.handlers.push(Box::new(future.into_future()))
    }

    /// Returns a handle to the internal event loop. This is a sort of escape hatch that allows you
    /// to take more control over what runs on the reactor using `tokio`. This can be used for
    /// sharing this reactor with some elements of other libraries.
    pub fn inner_handle(&self) -> Handle {
        self.inner.handle()
    }

    /// Consumes all registered handlers and futures, and runs them. When using
    /// `register_client_with_handler`, this will block indefinitely (until failure occurs) as it
    /// will simply continue to process new, incoming messages for each client that was registered.
    ///
    /// # Example
    /// ```no_run
    /// # extern crate irc;
    /// # use std::default::Default;
    /// # use irc::client::prelude::*;
    /// # use irc::error;
    /// # fn main() {
    /// # let config = Config::default();
    /// let mut reactor = IrcReactor::new().unwrap();
    /// let client = reactor.prepare_client_and_connect(&config).unwrap();
    /// reactor.register_client_with_handler(client, process_msg);
    /// reactor.run().unwrap();
    /// # }
    /// # fn process_msg(client: &IrcClient, message: Message) -> error::Result<()> { Ok(()) }
    /// ```
    pub fn run(&mut self) -> error::Result<()> {
        let mut handlers = Vec::new();
        while let Some(handler) = self.handlers.pop() {
            handlers.push(handler);
        }
        self.inner.run(future::join_all(handlers).map(|_| ()))
    }
}
<commit_msg>spawn handler futures on the reactor<commit_after>//! A system for creating and managing IRC client connections.
//!
//! This API provides the ability to create and manage multiple IRC clients that can run on the same
//! thread through the use of a shared event loop. It can also be used to encapsulate the dependency
//! on `tokio` and `futures` in the use of `IrcClient::new_future`. This means that knowledge of
//! those libraries should be unnecessary for the average user. Nevertheless, this API also provides
//! some escape hatches that let advanced users take further advantage of these dependencies.
//!
//! # Example
//! ```no_run
//! # extern crate irc;
//! # use std::default::Default;
//! use irc::client::prelude::*;
//! use irc::error;
//!
//! fn main() {
//!   let config = Config::default();
//!   let mut reactor = IrcReactor::new().unwrap();
//!   let client = reactor.prepare_client_and_connect(&config).unwrap();
//!   reactor.register_client_with_handler(client, process_msg);
//!   reactor.run().unwrap();
//! }
//! # fn process_msg(client: &IrcClient, message: Message) -> error::Result<()> { Ok(()) }
//! ```

use futures::{Future, IntoFuture, Stream};
use futures::future;
use tokio_core::reactor::{Core, Handle};

use client::data::Config;
use client::{IrcClient, IrcClientFuture, PackedIrcClient, Client};
use error;
use proto::Message;

/// A thin wrapper over an event loop.
///
/// An IRC reactor is used to create new IRC clients and to drive the management of all connected
/// clients as the application runs. It can be used to run multiple clients on the same thread, as
/// well as to get better control over error management in an IRC client.
///
/// For a full example usage, see [`irc::client::reactor`](./index.html).
pub struct IrcReactor {
    inner: Core,
    handlers: Vec<Box<Future<Item = (), Error = error::IrcError>>>,
}

impl IrcReactor {
    /// Creates a new reactor.
    pub fn new() -> error::Result<IrcReactor> {
        Ok(IrcReactor {
            inner: Core::new()?,
            handlers: Vec::new(),
        })
    }

    /// Creates a representation of an IRC client that has not yet attempted to connect. In
    /// particular, this representation is as a `Future` that when run will produce a connected
    /// [`IrcClient`](../struct.IrcClient.html).
    ///
    /// # Example
    /// ```no_run
    /// # extern crate irc;
    /// # use std::default::Default;
    /// # use irc::client::prelude::*;
    /// # fn main() {
    /// # let config = Config::default();
    /// let future_client = IrcReactor::new().and_then(|mut reactor| {
    ///     reactor.prepare_client(&config)
    /// });
    /// # }
    /// ```
    pub fn prepare_client<'a>(&mut self, config: &'a Config) -> error::Result<IrcClientFuture<'a>> {
        IrcClient::new_future(self.inner_handle(), config)
    }

    /// Runs an [`IrcClientFuture`](../struct.IrcClientFuture.html), such as one from
    /// `prepare_client` to completion, yielding an [`IrcClient`](../struct.IrcClient.html).
    ///
    /// # Example
    /// ```no_run
    /// # extern crate irc;
    /// # use std::default::Default;
    /// # use irc::client::prelude::*;
    /// # fn main() {
    /// # let config = Config::default();
    /// let client = IrcReactor::new().and_then(|mut reactor| {
    ///     reactor.prepare_client(&config).and_then(|future| {
    ///         reactor.connect_client(future)
    ///     })
    /// });
    /// # }
    /// ```
    pub fn connect_client(&mut self, future: IrcClientFuture) -> error::Result<IrcClient> {
        self.inner.run(future).map(|PackedIrcClient(client, future)| {
            self.register_future(future);
            client
        })
    }

    /// Creates a new [`IrcClient`](../struct.IrcClient.html) from the specified configuration,
    /// connecting immediately. This is guaranteed to be the composition of `prepare_client` and
    /// `connect_client`.
    ///
    /// # Example
    /// ```no_run
    /// # extern crate irc;
    /// # use std::default::Default;
    /// # use irc::client::prelude::*;
    /// # fn main() {
    /// # let config = Config::default();
    /// let client = IrcReactor::new().and_then(|mut reactor| {
    ///     reactor.prepare_client_and_connect(&config)
    /// });
    /// # }
    /// ```
    pub fn prepare_client_and_connect(&mut self, config: &Config) -> error::Result<IrcClient> {
        self.prepare_client(config).and_then(|future| self.connect_client(future))
    }

    /// Registers the given client with the specified message handler. The reactor will store this
    /// setup until the next call to run, where it will be used to process new messages over the
    /// connection indefinitely (or until failure). As registration is consumed by `run`, subsequent
    /// calls to run will require new registration.
    ///
    /// # Example
    /// ```no_run
    /// # extern crate irc;
    /// # use std::default::Default;
    /// # use irc::client::prelude::*;
    /// # fn main() {
    /// # let config = Config::default();
    /// let mut reactor = IrcReactor::new().unwrap();
    /// let client = reactor.prepare_client_and_connect(&config).unwrap();
    /// reactor.register_client_with_handler(client, |client, msg| {
    ///   // Message processing happens here.
    ///   Ok(())
    /// })
    /// # }
    /// ```
    pub fn register_client_with_handler<F, U>(
        &mut self, client: IrcClient, mut handler: F
    ) where F: FnMut(&IrcClient, Message) -> U + 'static,
            U: IntoFuture<Item = (), Error = error::IrcError> + 'static {
        let handle = self.inner.handle().clone();
        self.handlers.push(Box::new(client.stream().for_each(move |message| {
            handle.spawn(handler(&client, message).into_future().map_err(|_| (())));

            Ok(())
        })));
    }

    /// Registers an arbitrary future with this reactor. This is a sort of escape hatch that allows
    /// you to take more control over what runs on the reactor without requiring you to bring in
    /// additional knowledge about `tokio`. It is suspected that `register_client_with_handler` will
    /// be sufficient for most use cases.
    pub fn register_future<F>(
        &mut self, future: F
    ) where F: IntoFuture<Item = (), Error = error::IrcError> + 'static {
        self.handlers.push(Box::new(future.into_future()))
    }

    /// Returns a handle to the internal event loop. This is a sort of escape hatch that allows you
    /// to take more control over what runs on the reactor using `tokio`. This can be used for
    /// sharing this reactor with some elements of other libraries.
    pub fn inner_handle(&self) -> Handle {
        self.inner.handle()
    }

    /// Consumes all registered handlers and futures, and runs them. When using
    /// `register_client_with_handler`, this will block indefinitely (until failure occurs) as it
    /// will simply continue to process new, incoming messages for each client that was registered.
    ///
    /// # Example
    /// ```no_run
    /// # extern crate irc;
    /// # use std::default::Default;
    /// # use irc::client::prelude::*;
    /// # use irc::error;
    /// # fn main() {
    /// # let config = Config::default();
    /// let mut reactor = IrcReactor::new().unwrap();
    /// let client = reactor.prepare_client_and_connect(&config).unwrap();
    /// reactor.register_client_with_handler(client, process_msg);
    /// reactor.run().unwrap();
    /// # }
    /// # fn process_msg(client: &IrcClient, message: Message) -> error::Result<()> { Ok(()) }
    /// ```
    pub fn run(&mut self) -> error::Result<()> {
        let mut handlers = Vec::new();
        while let Some(handler) = self.handlers.pop() {
            handlers.push(handler);
        }
        self.inner.run(future::join_all(handlers).map(|_| ()))
    }
}
<|endoftext|>


<commit_before>//! Licensed under the Apache License, Version 2.0
//! http://www.apache.org/licenses/LICENSE-2.0 or the MIT license
//! http://opensource.org/licenses/MIT, at your
//! option. This file may not be copied, modified, or distributed
//! except according to those terms.

use std::num::Saturating;

/// An iterator like `.map(|elt| elt.clone())`
///
/// Iterator element type is `A`, if `I: Iterator<&A>`
pub struct Clones {
    iter: I
}

impl Clones
{
    pub fn new(iter: I) -> Clones {
        Clones{iter: iter}
    }
}

impl<'a, A: Clone, I: Iterator<&'a A>>
Iterator for Clones
{
    #[inline]
    fn next(&mut self) -> Option
    {
        self.iter.next().map(|elt| elt.clone())
    }

    fn size_hint(&self) -> (uint, Option<uint>)
    {
        self.iter.size_hint()
    }
}

impl<'a, A: Clone, I: DoubleEndedIterator<&'a A>>
DoubleEndedIterator for Clones
{
    #[inline]
    fn next_back(&mut self) -> Option
    {
        self.iter.next_back().map(|elt| elt.clone())
    }
}

impl<'a, A: Clone, I: ExactSize<&'a A>> ExactSize for Clones { }

/// Alternate elements from two iterators until both
/// are run out
///
/// Iterator element type is `A` if `I: Iterator<A>`
#[deriving(Clone)]
pub struct Interleave {
    a: I,
    b: J,
    flag: bool,
}

impl Interleave {
    ///
    pub fn new(a: I, b: J) -> Interleave {
        Interleave{a: a, b: b, flag: false}
    }
}

impl<A, I: Iterator, J: Iterator> Iterator for Interleave {
    #[inline]
    fn next(&mut self) -> Option {
        self.flag = !self.flag;
        if self.flag {
            match self.a.next() {
                None => self.b.next(),
                r => r,
            }
        } else {
            match self.b.next() {
                None => self.a.next(),
                r => r,
            }
        }
    }
}

/// Clonable iterator adaptor to map elementwise
/// from `Iterator<A>` to `Iterator<B>`
///
/// Created with `.fn_map(..)` on an iterator
///
/// Iterator element type is `B`
pub struct FnMap {
    map: fn(A) -> B,
    iter: I,
}

impl FnMap
{
    pub fn new(iter: I, map: fn(A) -> B) -> FnMap {
        FnMap{iter: iter, map: map}
    }
}

impl<A, B, I: Iterator> Iterator for FnMap
{
    #[inline]
    fn next(&mut self) -> Option {
        self.iter.next().map(|a| (self.map)(a))
    }

    fn size_hint(&self) -> (uint, Option<uint>) {
        self.iter.size_hint()
    }
}

impl<A, B, I: DoubleEndedIterator> DoubleEndedIterator
for FnMap
{
    #[inline]
    fn next_back(&mut self) -> Option {
        self.iter.next_back().map(|a| (self.map)(a))
    }
}

impl Clone for FnMap
{
    fn clone(&self) -> FnMap {
        FnMap::new(self.iter.clone(), self.map)
    }
}

/// An iterator adaptor that allows putting back a single
/// item to the front of the iterator.
///
/// Iterator element type is `A`
#[deriving(Clone)]
pub struct PutBack {
    top: Option,
    iter: I
}

impl PutBack {
    /// Iterator element type is `A`
    #[inline]
    pub fn new(it: I) -> PutBack {
        PutBack{top: None, iter: it}
    }

    /// Put back a single value to the front of the iterator.
    ///
    /// If a value is already in the put back slot, it is overwritten.
    #[inline]
    pub fn put_back(&mut self, x: A) {
        self.top = Some(x)
    }
}

impl<A, I: Iterator> Iterator for PutBack {
    #[inline]
    fn next(&mut self) -> Option {
        match self.top.take() {
            None => self.iter.next(),
            top => top,
        }
    }
    #[inline]
    fn size_hint(&self) -> (uint, Option<uint>) {
        let (lo, hi) = self.iter.size_hint();
        match self.top {
            Some(_) => (lo.saturating_add(1), hi.and_then(|x| x.checked_add(&1))),
            None => (lo, hi)
        }
    }
}

/// An iterator adaptor that iterates over the cartesian product of
/// the element sets of two iterators `I` and `J`.
///
/// Iterator element type is `(A, B)` if `I: Iterator<A>` and `J: Iterator<B>`
#[deriving(Clone)]
pub struct Product {
    a: I,
    a_cur: Option,
    b: J,
    b_orig: J,
}

impl<A: Clone, B, I: Iterator, J: Clone + Iterator>
    Product
{
    /// Create a new cartesian product iterator
    ///
    /// Iterator element type is `(A, B)` if `I: Iterator<A>` and `J: Iterator<B>`
    pub fn new(i: I, j: J) -> Product
    {
        let mut i = i;
        Product{a_cur: i.next(), a: i, b: j.clone(), b_orig: j}
    }
}

impl<A: Clone, I: Iterator, B, J: Clone + Iterator>
Iterator<(A, B)> for Product
{
    fn next(&mut self) -> Option<(A, B)>
    {
        let elt_b = match self.b.next() {
            None => {
                self.b = self.b_orig.clone();
                match self.b.next() {
                    None => return None,
                    Some(x) => {
                        self.a_cur = self.a.next();
                        x
                    }
                }
            }
            Some(x) => x
        };
        match self.a_cur {
            None => None,
            Some(ref a) => {
                Some((a.clone(), elt_b))
            }
        }
    }

    fn size_hint(&self) -> (uint, Option<uint>)
    {
        let (a, ah) = self.a.size_hint();
        let (b, bh) = self.b.size_hint();
        let (bo, boh) = self.b_orig.size_hint();

        // Compute a * bo + b for both lower and upper bound
        let low = a.checked_mul(&bo)
                    .and_then(|x| x.checked_add(&b))
                    .unwrap_or(::std::uint::MAX);
        let high = ah.and_then(|x| boh.and_then(|y| x.checked_mul(&y)))
                     .and_then(|x| bh.and_then(|y| x.checked_add(&y)));
        (low, high)
    }
}

<commit_msg>Minisimplification of PutBack<commit_after>//! Licensed under the Apache License, Version 2.0
//! http://www.apache.org/licenses/LICENSE-2.0 or the MIT license
//! http://opensource.org/licenses/MIT, at your
//! option. This file may not be copied, modified, or distributed
//! except according to those terms.

use std::mem;
use std::num::Saturating;

/// An iterator like `.map(|elt| elt.clone())`
///
/// Iterator element type is `A`, if `I: Iterator<&A>`
pub struct Clones {
    iter: I
}

impl Clones
{
    pub fn new(iter: I) -> Clones {
        Clones{iter: iter}
    }
}

impl<'a, A: Clone, I: Iterator<&'a A>>
Iterator for Clones
{
    #[inline]
    fn next(&mut self) -> Option
    {
        self.iter.next().map(|elt| elt.clone())
    }

    fn size_hint(&self) -> (uint, Option<uint>)
    {
        self.iter.size_hint()
    }
}

impl<'a, A: Clone, I: DoubleEndedIterator<&'a A>>
DoubleEndedIterator for Clones
{
    #[inline]
    fn next_back(&mut self) -> Option
    {
        self.iter.next_back().map(|elt| elt.clone())
    }
}

impl<'a, A: Clone, I: ExactSize<&'a A>> ExactSize for Clones { }

/// Alternate elements from two iterators until both
/// are run out
///
/// Iterator element type is `A` if `I: Iterator<A>`
#[deriving(Clone)]
pub struct Interleave {
    a: I,
    b: J,
    flag: bool,
}

impl Interleave {
    ///
    pub fn new(a: I, b: J) -> Interleave {
        Interleave{a: a, b: b, flag: false}
    }
}

impl<A, I: Iterator, J: Iterator> Iterator for Interleave {
    #[inline]
    fn next(&mut self) -> Option {
        self.flag = !self.flag;
        if self.flag {
            match self.a.next() {
                None => self.b.next(),
                r => r,
            }
        } else {
            match self.b.next() {
                None => self.a.next(),
                r => r,
            }
        }
    }
}

/// Clonable iterator adaptor to map elementwise
/// from `Iterator<A>` to `Iterator<B>`
///
/// Created with `.fn_map(..)` on an iterator
///
/// Iterator element type is `B`
pub struct FnMap {
    map: fn(A) -> B,
    iter: I,
}

impl FnMap
{
    pub fn new(iter: I, map: fn(A) -> B) -> FnMap {
        FnMap{iter: iter, map: map}
    }
}

impl<A, B, I: Iterator> Iterator for FnMap
{
    #[inline]
    fn next(&mut self) -> Option {
        self.iter.next().map(|a| (self.map)(a))
    }

    fn size_hint(&self) -> (uint, Option<uint>) {
        self.iter.size_hint()
    }
}

impl<A, B, I: DoubleEndedIterator> DoubleEndedIterator
for FnMap
{
    #[inline]
    fn next_back(&mut self) -> Option {
        self.iter.next_back().map(|a| (self.map)(a))
    }
}

impl Clone for FnMap
{
    fn clone(&self) -> FnMap {
        FnMap::new(self.iter.clone(), self.map)
    }
}

/// An iterator adaptor that allows putting back a single
/// item to the front of the iterator.
///
/// Iterator element type is `A`
#[deriving(Clone)]
pub struct PutBack {
    top: Option,
    iter: I
}

impl PutBack {
    /// Iterator element type is `A`
    #[inline]
    pub fn new(it: I) -> PutBack {
        PutBack{top: None, iter: it}
    }

    /// Put back a single value to the front of the iterator.
    ///
    /// If a value is already in the put back slot, it is overwritten.
    #[inline]
    pub fn put_back(&mut self, x: A) {
        self.top = Some(x)
    }
}

impl<A, I: Iterator> Iterator for PutBack {
    #[inline]
    fn next(&mut self) -> Option {
        match self.top {
            None => self.iter.next(),
            ref mut some => mem::replace(some, None)
        }
    }
    #[inline]
    fn size_hint(&self) -> (uint, Option<uint>) {
        let (lo, hi) = self.iter.size_hint();
        match self.top {
            Some(_) => (lo.saturating_add(1), hi.and_then(|x| x.checked_add(&1))),
            None => (lo, hi)
        }
    }
}

/// An iterator adaptor that iterates over the cartesian product of
/// the element sets of two iterators `I` and `J`.
///
/// Iterator element type is `(A, B)` if `I: Iterator<A>` and `J: Iterator<B>`
#[deriving(Clone)]
pub struct Product {
    a: I,
    a_cur: Option,
    b: J,
    b_orig: J,
}

impl<A: Clone, B, I: Iterator, J: Clone + Iterator>
    Product
{
    /// Create a new cartesian product iterator
    ///
    /// Iterator element type is `(A, B)` if `I: Iterator<A>` and `J: Iterator<B>`
    pub fn new(i: I, j: J) -> Product
    {
        let mut i = i;
        Product{a_cur: i.next(), a: i, b: j.clone(), b_orig: j}
    }
}

impl<A: Clone, I: Iterator, B, J: Clone + Iterator>
Iterator<(A, B)> for Product
{
    fn next(&mut self) -> Option<(A, B)>
    {
        let elt_b = match self.b.next() {
            None => {
                self.b = self.b_orig.clone();
                match self.b.next() {
                    None => return None,
                    Some(x) => {
                        self.a_cur = self.a.next();
                        x
                    }
                }
            }
            Some(x) => x
        };
        match self.a_cur {
            None => None,
            Some(ref a) => {
                Some((a.clone(), elt_b))
            }
        }
    }

    fn size_hint(&self) -> (uint, Option<uint>)
    {
        let (a, ah) = self.a.size_hint();
        let (b, bh) = self.b.size_hint();
        let (bo, boh) = self.b_orig.size_hint();

        // Compute a * bo + b for both lower and upper bound
        let low = a.checked_mul(&bo)
                    .and_then(|x| x.checked_add(&b))
                    .unwrap_or(::std::uint::MAX);
        let high = ah.and_then(|x| boh.and_then(|y| x.checked_mul(&y)))
                     .and_then(|x| bh.and_then(|y| x.checked_add(&y)));
        (low, high)
    }
}

<|endoftext|>


<commit_before>//!

use rand;

use super::super::visit::{Visitable, BucketCollector};
use super::super::structure::{Program, Sensor, Command, Expression, Condition};
use super::super::copy;
use super::super::copy::Copyable;

#[derive(Copy, Clone)]
enum NodeType {
    Program,
    Expression,
    Condition,
    Sensor,
    Command
}

pub fn cross_over<H: rand::Rng + Sized>(a: &Program, b: &Program, rng: &mut H) -> (Box<Program>, Box<Program>) {
    let mut a_nodes = BucketCollector::new();
    let mut b_nodes = BucketCollector::new();

    a.visit(&mut a_nodes);
    b.visit(&mut b_nodes);

    let nodes = get_counts(&a_nodes).into_iter()
        .zip(get_counts(&b_nodes).into_iter())
        .filter(|&((_, ac), (_, bc))| ac > 0 && bc > 0)
        .map(|((t, _), (_, _))| t)
        .collect::<Vec<NodeType>>();

    let picked_type = choose(nodes, rng);

    match picked_type {
        NodeType::Program    => cross_over_program(a, b, choose(a_nodes.programs, rng), choose(b_nodes.programs, rng), rng),
        NodeType::Expression => cross_over_expression(a, b, choose(a_nodes.expressions, rng), choose(b_nodes.expressions, rng), rng),
        NodeType::Condition  => cross_over_condition(a, b, choose(a_nodes.conditions, rng), choose(b_nodes.conditions, rng), rng),
        NodeType::Sensor     => cross_over_sensor(a, b, choose(a_nodes.sensors, rng), choose(b_nodes.sensors, rng), rng),
        NodeType::Command    => cross_over_command(a, b, choose(a_nodes.commands, rng), choose(b_nodes.commands, rng), rng)
    }
}

fn cross_over_program<'a, H: rand::Rng+Sized>(a: &Program, b: &Program, an: &Program, bn: &Program, rng: &mut H) -> (Box<Program>, Box<Program>) {
    (a.copy(&copy::CopyReplaceProgram { to_replace: an, replace_with: bn }),
     b.copy(&copy::CopyReplaceProgram { to_replace: bn, replace_with: an }))
}

fn cross_over_expression<'a, H: rand::Rng+Sized>(a: &Program, b: &Program, an: &Expression, bn: &Expression, rng: &mut H) -> (Box<Program>, Box<Program>) {
    (a.copy(&copy::CopyReplaceExpression { to_replace: an, replace_with: bn }),
     b.copy(&copy::CopyReplaceExpression { to_replace: bn, replace_with: an }))
}

fn cross_over_condition<'a, H: rand::Rng+Sized>(a: &Program, b: &Program, an: &Condition, bn: &Condition, rng: &mut H) -> (Box<Program>, Box<Program>) {
    (a.copy(&copy::CopyReplaceCondition { to_replace: an, replace_with: bn }),
     b.copy(&copy::CopyReplaceCondition { to_replace: bn, replace_with: an }))
}

fn cross_over_sensor<'a, H: rand::Rng+Sized>(a: &Program, b: &Program, an: &Sensor, bn: &Sensor, rng: &mut H) -> (Box<Program>, Box<Program>) {
    (a.copy(&copy::CopyReplaceSensor { to_replace: an, replace_with: bn }),
     b.copy(&copy::CopyReplaceSensor { to_replace: bn, replace_with: an }))
}

fn cross_over_command<'a, H: rand::Rng+Sized>(a: &Program, b: &Program, an: &Command, bn: &Command, rng: &mut H) -> (Box<Program>, Box<Program>) {
    (a.copy(&copy::CopyReplaceCommand { to_replace: an, replace_with: bn }),
     b.copy(&copy::CopyReplaceCommand { to_replace: bn, replace_with: an }))
}

fn get_counts(coll: &BucketCollector) -> Vec<(NodeType, usize)> {
    vec![(NodeType::Program,    coll.programs.len()),
         (NodeType::Expression, coll.expressions.len()),
         (NodeType::Condition,  coll.conditions.len()),
         (NodeType::Sensor,     coll.sensors.len()),
         (NodeType::Command,    coll.commands.len())]
}

fn choose<T: Copy, H: rand::Rng + Sized>(vec: Vec<T>, rng: &mut H) -> T {
    let i = rng.gen_range(0, vec.len());
    vec[i]
}
<commit_msg>Get rid of useless type annotations in crossover<commit_after>//!

use rand;

use super::super::visit::{Visitable, BucketCollector};
use super::super::structure::{Program, Sensor, Command, Expression, Condition};
use super::super::copy;
use super::super::copy::Copyable;

#[derive(Copy, Clone)]
enum NodeType {
    Program,
    Expression,
    Condition,
    Sensor,
    Command
}

pub fn cross_over<H: rand::Rng + Sized>(a: &Program, b: &Program, rng: &mut H) -> (Box<Program>, Box<Program>) {
    let mut a_nodes = BucketCollector::new();
    let mut b_nodes = BucketCollector::new();

    a.visit(&mut a_nodes);
    b.visit(&mut b_nodes);

    let nodes = get_counts(&a_nodes).into_iter()
        .zip(get_counts(&b_nodes).into_iter())
        .filter(|&((_, ac), (_, bc))| ac > 0 && bc > 0)
        .map(|((t, _), (_, _))| t)
        .collect::<Vec<NodeType>>();

    let picked_type = choose(nodes, rng);

    match picked_type {
        NodeType::Program    => cross_over_program(a, b, choose(a_nodes.programs, rng), choose(b_nodes.programs, rng)),
        NodeType::Expression => cross_over_expression(a, b, choose(a_nodes.expressions, rng), choose(b_nodes.expressions, rng)),
        NodeType::Condition  => cross_over_condition(a, b, choose(a_nodes.conditions, rng), choose(b_nodes.conditions, rng)),
        NodeType::Sensor     => cross_over_sensor(a, b, choose(a_nodes.sensors, rng), choose(b_nodes.sensors, rng)),
        NodeType::Command    => cross_over_command(a, b, choose(a_nodes.commands, rng), choose(b_nodes.commands, rng))
    }
}

fn cross_over_program(a: &Program, b: &Program, an: &Program, bn: &Program) -> (Box<Program>, Box<Program>) {
    (a.copy(&copy::CopyReplaceProgram { to_replace: an, replace_with: bn }),
     b.copy(&copy::CopyReplaceProgram { to_replace: bn, replace_with: an }))
}

fn cross_over_expression(a: &Program, b: &Program, an: &Expression, bn: &Expression) -> (Box<Program>, Box<Program>) {
    (a.copy(&copy::CopyReplaceExpression { to_replace: an, replace_with: bn }),
     b.copy(&copy::CopyReplaceExpression { to_replace: bn, replace_with: an }))
}

fn cross_over_condition(a: &Program, b: &Program, an: &Condition, bn: &Condition) -> (Box<Program>, Box<Program>) {
    (a.copy(&copy::CopyReplaceCondition { to_replace: an, replace_with: bn }),
     b.copy(&copy::CopyReplaceCondition { to_replace: bn, replace_with: an }))
}

fn cross_over_sensor(a: &Program, b: &Program, an: &Sensor, bn: &Sensor) -> (Box<Program>, Box<Program>) {
    (a.copy(&copy::CopyReplaceSensor { to_replace: an, replace_with: bn }),
     b.copy(&copy::CopyReplaceSensor { to_replace: bn, replace_with: an }))
}

fn cross_over_command(a: &Program, b: &Program, an: &Command, bn: &Command) -> (Box<Program>, Box<Program>) {
    (a.copy(&copy::CopyReplaceCommand { to_replace: an, replace_with: bn }),
     b.copy(&copy::CopyReplaceCommand { to_replace: bn, replace_with: an }))
}

fn get_counts(coll: &BucketCollector) -> Vec<(NodeType, usize)> {
    vec![(NodeType::Program,    coll.programs.len()),
         (NodeType::Expression, coll.expressions.len()),
         (NodeType::Condition,  coll.conditions.len()),
         (NodeType::Sensor,     coll.sensors.len()),
         (NodeType::Command,    coll.commands.len())]
}

fn choose<T: Copy, H: rand::Rng + Sized>(vec: Vec<T>, rng: &mut H) -> T {
    let i = rng.gen_range(0, vec.len());
    vec[i]
}
<|endoftext|>


<commit_before>// Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! Handles argument parsing.
//!
//! # Example
//!
//! ```
//! # use crosvm::argument::{Argument, Error, print_help, set_arguments};
//! # let args: std::slice::Iter<String> = [].iter();
//! let arguments = &[
//!     Argument::positional("FILES", "files to operate on"),
//!     Argument::short_value('p', "program", "PROGRAM", "Program to apply to each file"),
//!     Argument::short_value('c', "cpus", "N", "Number of CPUs to use. (default: 1)"),
//!     Argument::flag("unmount", "Unmount the root"),
//!     Argument::short_flag('h', "help", "Print help message."),
//! ];
//!
//! let match_res = set_arguments(args, arguments, |name, value| {
//!     match name {
//!         "" => println!("positional arg! {}", value.unwrap()),
//!         "program" => println!("gonna use program {}", value.unwrap()),
//!         "cpus" => {
//!             let v: u32 = value.unwrap().parse().map_err(|_| {
//!                 Error::InvalidValue {
//!                     value: value.unwrap().to_owned(),
//!                     expected: "this value for `cpus` needs to be integer",
//!                 }
//!             })?;
//!         }
//!         "unmount" => println!("gonna unmount"),
//!         "help" => return Err(Error::PrintHelp),
//!         _ => unreachable!(),
//!     }
//!     unreachable!();
//! });
//!
//! match match_res {
//!     Ok(_) => println!("running with settings"),
//!     Err(Error::PrintHelp) => print_help("best_program", "FILES", arguments),
//!     Err(e) => println!("{}", e),
//! }
//! ```

use std::fmt::{self, Display};
use std::result;

/// An error with argument parsing.
#[derive(Debug)]
pub enum Error {
    /// There was a syntax error with the argument.
    Syntax(String),
    /// The argumen's name is unused.
    UnknownArgument(String),
    /// The argument was required.
    ExpectedArgument(String),
    /// The argument's given value is invalid.
    InvalidValue {
        value: String,
        expected: &'static str,
    },
    /// The argument was already given and none more are expected.
    TooManyArguments(String),
    /// The argument expects a value.
    ExpectedValue(String),
    /// The help information was requested
    PrintHelp,
}

impl Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::Error::*;

        match self {
            Syntax(s) => write!(f, "syntax error: {}", s),
            UnknownArgument(s) => write!(f, "unknown argument: {}", s),
            ExpectedArgument(s) => write!(f, "expected argument: {}", s),
            InvalidValue { value, expected } => {
                write!(f, "invalid value {:?}: {}", value, expected)
            }
            TooManyArguments(s) => write!(f, "too many arguments: {}", s),
            ExpectedValue(s) => write!(f, "expected parameter value: {}", s),
            PrintHelp => write!(f, "help was requested"),
        }
    }
}

/// Result of a argument parsing.
pub type Result<T> = result::Result<T, Error>;

/// Information about an argument expected from the command line.
///
/// # Examples
///
/// To indicate a flag style argument:
///
/// ```
/// # use crosvm::argument::Argument;
/// Argument::short_flag('f', "flag", "enable awesome mode");
/// ```
///
/// To indicate a parameter style argument that expects a value:
///
/// ```
/// # use crosvm::argument::Argument;
/// // "VALUE" and "NETMASK" are placeholder values displayed in the help message for these
/// // arguments.
/// Argument::short_value('v', "val", "VALUE", "how much do you value this usage information");
/// Argument::value("netmask", "NETMASK", "hides your netface");
/// ```
///
/// To indicate an argument with no short version:
///
/// ```
/// # use crosvm::argument::Argument;
/// Argument::flag("verbose", "this option is hard to type quickly");
/// ```
///
/// To indicate a positional argument:
///
/// ```
/// # use crosvm::argument::Argument;
/// Argument::positional("VALUES", "these are positional arguments");
/// ```
#[derive(Default)]
pub struct Argument {
    /// The name of the value to display in the usage information. Use None to indicate that there
    /// is no value expected for this argument.
    pub value: Option<&'static str>,
    /// Optional single character shortened argument name.
    pub short: Option<char>,
    /// The long name of this argument.
    pub long: &'static str,
    /// Helpfuly usage information for this argument to display to the user.
    pub help: &'static str,
}

impl Argument {
    pub fn positional(value: &'static str, help: &'static str) -> Argument {
        Argument {
            value: Some(value),
            long: "",
            help,
            ..Default::default()
        }
    }

    pub fn value(long: &'static str, value: &'static str, help: &'static str) -> Argument {
        Argument {
            value: Some(value),
            long,
            help,
            ..Default::default()
        }
    }

    pub fn short_value(
        short: char,
        long: &'static str,
        value: &'static str,
        help: &'static str,
    ) -> Argument {
        Argument {
            value: Some(value),
            short: Some(short),
            long,
            help,
        }
    }

    pub fn flag(long: &'static str, help: &'static str) -> Argument {
        Argument {
            long,
            help,
            ..Default::default()
        }
    }

    pub fn short_flag(short: char, long: &'static str, help: &'static str) -> Argument {
        Argument {
            short: Some(short),
            long,
            help,
            ..Default::default()
        }
    }
}

fn parse_arguments<I, R, F>(args: I, mut f: F) -> Result<()>
where
    I: Iterator<Item = R>,
    R: AsRef<str>,
    F: FnMut(&str, Option<&str>) -> Result<()>,
{
    enum State {
        // Initial state at the start and after finishing a single argument/value.
        Top,
        // The remaining arguments are all positional.
        Positional,
        // The next string is the value for the argument `name`.
        Value { name: String },
    }
    let mut s = State::Top;
    for arg in args {
        let arg = arg.as_ref();
        s = match s {
            State::Top => {
                if arg == "--" {
                    State::Positional
                } else if arg.starts_with("--") {
                    let param = arg.trim_start_matches('-');
                    if param.contains('=') {
                        let mut iter = param.splitn(2, '=');
                        let name = iter.next().unwrap();
                        let value = iter.next().unwrap();
                        if name.is_empty() {
                            return Err(Error::Syntax(
                                "expected parameter name before `=`".to_owned(),
                            ));
                        }
                        if value.is_empty() {
                            return Err(Error::Syntax(
                                "expected parameter value after `=`".to_owned(),
                            ));
                        }
                        f(name, Some(value))?;
                        State::Top
                    } else if let Err(e) = f(param, None) {
                        if let Error::ExpectedValue(_) = e {
                            State::Value {
                                name: param.to_owned(),
                            }
                        } else {
                            return Err(e);
                        }
                    } else {
                        State::Top
                    }
                } else if arg.starts_with('-') {
                    if arg.len() == 1 {
                        return Err(Error::Syntax(
                            "expected argument short name after `-`".to_owned(),
                        ));
                    }
                    let name = &arg[1..2];
                    let value = if arg.len() > 2 { Some(&arg[2..]) } else { None };
                    if let Err(e) = f(name, value) {
                        if let Error::ExpectedValue(_) = e {
                            State::Value {
                                name: name.to_owned(),
                            }
                        } else {
                            return Err(e);
                        }
                    } else {
                        State::Top
                    }
                } else {
                    f("", Some(&arg))?;
                    State::Positional
                }
            }
            State::Positional => {
                f("", Some(&arg))?;
                State::Positional
            }
            State::Value { name } => {
                f(&name, Some(&arg))?;
                State::Top
            }
        };
    }
    Ok(())
}

/// Parses the given `args` against the list of know arguments `arg_list` and calls `f` with each
/// present argument and value if required.
///
/// This function guarantees that only valid long argument names from `arg_list` are sent to the
/// callback `f`. It is also guaranteed that if an arg requires a value (i.e.
/// `arg.value.is_some()`), the value will be `Some` in the callbacks arguments. If the callback
/// returns `Err`, this function will end parsing and return that `Err`.
///
/// See the [module level](index.html) example for a usage example.
pub fn set_arguments<I, R, F>(args: I, arg_list: &[Argument], mut f: F) -> Result<()>
where
    I: Iterator<Item = R>,
    R: AsRef<str>,
    F: FnMut(&str, Option<&str>) -> Result<()>,
{
    parse_arguments(args, |name, value| {
        let mut matches = None;
        for arg in arg_list {
            if let Some(short) = arg.short {
                if name.len() == 1 && name.starts_with(short) {
                    if value.is_some() != arg.value.is_some() {
                        return Err(Error::ExpectedValue(short.to_string()));
                    }
                    matches = Some(arg.long);
                }
            }
            if matches.is_none() && arg.long == name {
                if value.is_some() != arg.value.is_some() {
                    return Err(Error::ExpectedValue(arg.long.to_owned()));
                }
                matches = Some(arg.long);
            }
        }
        match matches {
            Some(long) => f(long, value),
            None => Err(Error::UnknownArgument(name.to_owned())),
        }
    })
}

/// Prints command line usage information to stdout.
///
/// Usage information is printed according to the help fields in `args` with a leading usage line.
/// The usage line is of the format "`program_name` [ARGUMENTS] `required_arg`".
pub fn print_help(program_name: &str, required_arg: &str, args: &[Argument]) {
    println!(
        "Usage: {} {}{}\n",
        program_name,
        if args.is_empty() { "" } else { "[ARGUMENTS] " },
        required_arg
    );
    if args.is_empty() {
        return;
    }
    println!("Argument{}:", if args.len() > 1 { "s" } else { "" });
    for arg in args {
        match arg.short {
            Some(s) => print!(" -{}, ", s),
            None => print!("     "),
        }
        if arg.long.is_empty() {
            print!("  ");
        } else {
            print!("--");
        }
        print!("{:<12}", arg.long);
        if let Some(v) = arg.value {
            if arg.long.is_empty() {
                print!(" ");
            } else {
                print!("=");
            }
            print!("{:<10}", v);
        } else {
            print!("{:<11}", "");
        }
        println!("{}", arg.help);
    }
}

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

    #[test]
    fn request_help() {
        let arguments = [Argument::short_flag('h', "help", "Print help message.")];

        let match_res = set_arguments(["-h"].iter(), &arguments[..], |name, _| {
            match name {
                "help" => return Err(Error::PrintHelp),
                _ => unreachable!(),
            };
        });
        match match_res {
            Err(Error::PrintHelp) => {}
            _ => unreachable!(),
        }
    }

    #[test]
    fn mixed_args() {
        let arguments = [
            Argument::positional("FILES", "files to operate on"),
            Argument::short_value('p', "program", "PROGRAM", "Program to apply to each file"),
            Argument::short_value('c', "cpus", "N", "Number of CPUs to use. (default: 1)"),
            Argument::flag("unmount", "Unmount the root"),
            Argument::short_flag('h', "help", "Print help message."),
        ];

        let mut unmount = false;
        let match_res = set_arguments(
            ["--cpus", "3", "--program", "hello", "--unmount", "file"].iter(),
            &arguments[..],
            |name, value| {
                match name {
                    "" => assert_eq!(value.unwrap(), "file"),
                    "program" => assert_eq!(value.unwrap(), "hello"),
                    "cpus" => {
                        let c: u32 = value.unwrap().parse().map_err(|_| Error::InvalidValue {
                            value: value.unwrap().to_owned(),
                            expected: "this value for `cpus` needs to be integer",
                        })?;
                        assert_eq!(c, 3);
                    }
                    "unmount" => unmount = true,
                    "help" => return Err(Error::PrintHelp),
                    _ => unreachable!(),
                };
                Ok(())
            },
        );
        assert!(match_res.is_ok());
        assert!(unmount);
    }

    #[test]
    fn name_value_pair() {
        let arguments = [Argument::short_value(
            'c',
            "cpus",
            "N",
            "Number of CPUs to use. (default: 1)",
        )];
        let match_res = set_arguments(
            ["-c", "5", "--cpus", "5", "-c5", "--cpus=5"].iter(),
            &arguments[..],
            |name, value| {
                assert_eq!(name, "cpus");
                assert_eq!(value, Some("5"));
                Ok(())
            },
        );
        assert!(match_res.is_ok());
    }
}
<commit_msg>Adds a flag_or_value command line argument option<commit_after>// Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! Handles argument parsing.
//!
//! # Example
//!
//! ```
//! # use crosvm::argument::{Argument, Error, print_help, set_arguments};
//! # let args: std::slice::Iter<String> = [].iter();
//! let arguments = &[
//!     Argument::positional("FILES", "files to operate on"),
//!     Argument::short_value('p', "program", "PROGRAM", "Program to apply to each file"),
//!     Argument::short_value('c', "cpus", "N", "Number of CPUs to use. (default: 1)"),
//!     Argument::flag("unmount", "Unmount the root"),
//!     Argument::short_flag('h', "help", "Print help message."),
//! ];
//!
//! let match_res = set_arguments(args, arguments, |name, value| {
//!     match name {
//!         "" => println!("positional arg! {}", value.unwrap()),
//!         "program" => println!("gonna use program {}", value.unwrap()),
//!         "cpus" => {
//!             let v: u32 = value.unwrap().parse().map_err(|_| {
//!                 Error::InvalidValue {
//!                     value: value.unwrap().to_owned(),
//!                     expected: "this value for `cpus` needs to be integer",
//!                 }
//!             })?;
//!         }
//!         "unmount" => println!("gonna unmount"),
//!         "help" => return Err(Error::PrintHelp),
//!         _ => unreachable!(),
//!     }
//!     unreachable!();
//! });
//!
//! match match_res {
//!     Ok(_) => println!("running with settings"),
//!     Err(Error::PrintHelp) => print_help("best_program", "FILES", arguments),
//!     Err(e) => println!("{}", e),
//! }
//! ```

use std::fmt::{self, Display};
use std::result;

/// An error with argument parsing.
#[derive(Debug)]
pub enum Error {
    /// There was a syntax error with the argument.
    Syntax(String),
    /// The argumen's name is unused.
    UnknownArgument(String),
    /// The argument was required.
    ExpectedArgument(String),
    /// The argument's given value is invalid.
    InvalidValue {
        value: String,
        expected: &'static str,
    },
    /// The argument was already given and none more are expected.
    TooManyArguments(String),
    /// The argument expects a value.
    ExpectedValue(String),
    /// The argument does not expect a value.
    UnexpectedValue(String),
    /// The help information was requested
    PrintHelp,
}

impl Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::Error::*;

        match self {
            Syntax(s) => write!(f, "syntax error: {}", s),
            UnknownArgument(s) => write!(f, "unknown argument: {}", s),
            ExpectedArgument(s) => write!(f, "expected argument: {}", s),
            InvalidValue { value, expected } => {
                write!(f, "invalid value {:?}: {}", value, expected)
            }
            TooManyArguments(s) => write!(f, "too many arguments: {}", s),
            ExpectedValue(s) => write!(f, "expected parameter value: {}", s),
            UnexpectedValue(s) => write!(f, "unexpected parameter value: {}", s),
            PrintHelp => write!(f, "help was requested"),
        }
    }
}

/// Result of a argument parsing.
pub type Result<T> = result::Result<T, Error>;

#[derive(Debug, PartialEq)]
pub enum ArgumentValueMode {
    /// Specifies that an argument requires a value and that an error should be generated if
    /// no value is provided during parsing.
    Required,

    /// Specifies that an argument does not allow a value and that an error should be returned
    /// if a value is provided during parsing.
    Disallowed,

    /// Specifies that an argument may have a value during parsing but is not required to.
    Optional,
}

/// Information about an argument expected from the command line.
///
/// # Examples
///
/// To indicate a flag style argument:
///
/// ```
/// # use crosvm::argument::Argument;
/// Argument::short_flag('f', "flag", "enable awesome mode");
/// ```
///
/// To indicate a parameter style argument that expects a value:
///
/// ```
/// # use crosvm::argument::Argument;
/// // "VALUE" and "NETMASK" are placeholder values displayed in the help message for these
/// // arguments.
/// Argument::short_value('v', "val", "VALUE", "how much do you value this usage information");
/// Argument::value("netmask", "NETMASK", "hides your netface");
/// ```
///
/// To indicate an argument with no short version:
///
/// ```
/// # use crosvm::argument::Argument;
/// Argument::flag("verbose", "this option is hard to type quickly");
/// ```
///
/// To indicate a positional argument:
///
/// ```
/// # use crosvm::argument::Argument;
/// Argument::positional("VALUES", "these are positional arguments");
/// ```
pub struct Argument {
    /// The name of the value to display in the usage information.
    pub value: Option<&'static str>,
    /// Specifies how values should be handled for this this argument.
    pub value_mode: ArgumentValueMode,
    /// Optional single character shortened argument name.
    pub short: Option<char>,
    /// The long name of this argument.
    pub long: &'static str,
    /// Helpfuly usage information for this argument to display to the user.
    pub help: &'static str,
}

impl Argument {
    pub fn positional(value: &'static str, help: &'static str) -> Argument {
        Argument {
            value: Some(value),
            value_mode: ArgumentValueMode::Required,
            short: None,
            long: "",
            help,
        }
    }

    pub fn value(long: &'static str, value: &'static str, help: &'static str) -> Argument {
        Argument {
            value: Some(value),
            value_mode: ArgumentValueMode::Required,
            short: None,
            long,
            help,
        }
    }

    pub fn short_value(
        short: char,
        long: &'static str,
        value: &'static str,
        help: &'static str,
    ) -> Argument {
        Argument {
            value: Some(value),
            value_mode: ArgumentValueMode::Required,
            short: Some(short),
            long,
            help,
        }
    }

    pub fn flag(long: &'static str, help: &'static str) -> Argument {
        Argument {
            value: None,
            value_mode: ArgumentValueMode::Disallowed,
            short: None,
            long,
            help,
        }
    }

    pub fn short_flag(short: char, long: &'static str, help: &'static str) -> Argument {
        Argument {
            value: None,
            value_mode: ArgumentValueMode::Disallowed,
            short: Some(short),
            long,
            help,
        }
    }

    pub fn flag_or_value(long: &'static str, value: &'static str, help: &'static str) -> Argument {
        Argument {
            value: Some(value),
            value_mode: ArgumentValueMode::Optional,
            short: None,
            long,
            help,
        }
    }
}

fn parse_arguments<I, R, F>(args: I, mut f: F) -> Result<()>
where
    I: Iterator<Item = R>,
    R: AsRef<str>,
    F: FnMut(&str, Option<&str>) -> Result<()>,
{
    enum State {
        // Initial state at the start and after finishing a single argument/value.
        Top,
        // The remaining arguments are all positional.
        Positional,
        // The next string is the value for the argument `name`.
        Value { name: String },
    }
    let mut s = State::Top;
    for arg in args {
        let arg = arg.as_ref();
        loop {
            let mut arg_consumed = true;
            s = match s {
                State::Top => {
                    if arg == "--" {
                        State::Positional
                    } else if arg.starts_with("--") {
                        let param = arg.trim_start_matches('-');
                        if param.contains('=') {
                            let mut iter = param.splitn(2, '=');
                            let name = iter.next().unwrap();
                            let value = iter.next().unwrap();
                            if name.is_empty() {
                                return Err(Error::Syntax(
                                    "expected parameter name before `=`".to_owned(),
                                ));
                            }
                            if value.is_empty() {
                                return Err(Error::Syntax(
                                    "expected parameter value after `=`".to_owned(),
                                ));
                            }
                            f(name, Some(value))?;
                            State::Top
                        } else {
                            State::Value {
                                name: param.to_owned(),
                            }
                        }
                    } else if arg.starts_with('-') {
                        if arg.len() == 1 {
                            return Err(Error::Syntax(
                                "expected argument short name after `-`".to_owned(),
                            ));
                        }
                        let name = &arg[1..2];
                        let value = if arg.len() > 2 { Some(&arg[2..]) } else { None };
                        if let Err(e) = f(name, value) {
                            if let Error::ExpectedValue(_) = e {
                                State::Value {
                                    name: name.to_owned(),
                                }
                            } else {
                                return Err(e);
                            }
                        } else {
                            State::Top
                        }
                    } else {
                        f("", Some(&arg))?;
                        State::Positional
                    }
                }
                State::Positional => {
                    f("", Some(&arg))?;
                    State::Positional
                }
                State::Value { name } => {
                    if arg.starts_with("-") {
                        arg_consumed = false;
                        f(&name, None)?;
                    } else {
                        if let Err(_) = f(&name, Some(&arg)) {
                            arg_consumed = false;
                            f(&name, None)?;
                        }
                    }
                    State::Top
                }
            };

            if arg_consumed {
                break;
            }
        }
    }
    Ok(())
}

/// Parses the given `args` against the list of know arguments `arg_list` and calls `f` with each
/// present argument and value if required.
///
/// This function guarantees that only valid long argument names from `arg_list` are sent to the
/// callback `f`. It is also guaranteed that if an arg requires a value (i.e.
/// `arg.value.is_some()`), the value will be `Some` in the callbacks arguments. If the callback
/// returns `Err`, this function will end parsing and return that `Err`.
///
/// See the [module level](index.html) example for a usage example.
pub fn set_arguments<I, R, F>(args: I, arg_list: &[Argument], mut f: F) -> Result<()>
where
    I: Iterator<Item = R>,
    R: AsRef<str>,
    F: FnMut(&str, Option<&str>) -> Result<()>,
{
    parse_arguments(args, |name, value| {
        let mut matches = None;
        for arg in arg_list {
            if let Some(short) = arg.short {
                if name.len() == 1 && name.starts_with(short) {
                    if value.is_some() != arg.value.is_some() {
                        return Err(Error::ExpectedValue(short.to_string()));
                    }
                    matches = Some(arg.long);
                }
            }
            if matches.is_none() && arg.long == name {
                if value.is_none() && arg.value_mode == ArgumentValueMode::Required {
                    return Err(Error::ExpectedValue(arg.long.to_owned()));
                }
                if value.is_some() && arg.value_mode == ArgumentValueMode::Disallowed {
                    return Err(Error::UnexpectedValue(arg.long.to_owned()));
                }
                matches = Some(arg.long);
            }
        }
        match matches {
            Some(long) => f(long, value),
            None => Err(Error::UnknownArgument(name.to_owned())),
        }
    })
}

/// Prints command line usage information to stdout.
///
/// Usage information is printed according to the help fields in `args` with a leading usage line.
/// The usage line is of the format "`program_name` [ARGUMENTS] `required_arg`".
pub fn print_help(program_name: &str, required_arg: &str, args: &[Argument]) {
    println!(
        "Usage: {} {}{}\n",
        program_name,
        if args.is_empty() { "" } else { "[ARGUMENTS] " },
        required_arg
    );
    if args.is_empty() {
        return;
    }
    println!("Argument{}:", if args.len() > 1 { "s" } else { "" });
    for arg in args {
        match arg.short {
            Some(s) => print!(" -{}, ", s),
            None => print!("     "),
        }
        if arg.long.is_empty() {
            print!("  ");
        } else {
            print!("--");
        }
        print!("{:<12}", arg.long);
        if let Some(v) = arg.value {
            if arg.long.is_empty() {
                print!(" ");
            } else {
                print!("=");
            }
            print!("{:<10}", v);
        } else {
            print!("{:<11}", "");
        }
        println!("{}", arg.help);
    }
}

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

    #[test]
    fn request_help() {
        let arguments = [Argument::short_flag('h', "help", "Print help message.")];

        let match_res = set_arguments(["-h"].iter(), &arguments[..], |name, _| {
            match name {
                "help" => return Err(Error::PrintHelp),
                _ => unreachable!(),
            };
        });
        match match_res {
            Err(Error::PrintHelp) => {}
            _ => unreachable!(),
        }
    }

    #[test]
    fn mixed_args() {
        let arguments = [
            Argument::positional("FILES", "files to operate on"),
            Argument::short_value('p', "program", "PROGRAM", "Program to apply to each file"),
            Argument::short_value('c', "cpus", "N", "Number of CPUs to use. (default: 1)"),
            Argument::flag("unmount", "Unmount the root"),
            Argument::short_flag('h', "help", "Print help message."),
        ];

        let mut unmount = false;
        let match_res = set_arguments(
            ["--cpus", "3", "--program", "hello", "--unmount", "file"].iter(),
            &arguments[..],
            |name, value| {
                match name {
                    "" => assert_eq!(value.unwrap(), "file"),
                    "program" => assert_eq!(value.unwrap(), "hello"),
                    "cpus" => {
                        let c: u32 = value.unwrap().parse().map_err(|_| Error::InvalidValue {
                            value: value.unwrap().to_owned(),
                            expected: "this value for `cpus` needs to be integer",
                        })?;
                        assert_eq!(c, 3);
                    }
                    "unmount" => unmount = true,
                    "help" => return Err(Error::PrintHelp),
                    _ => unreachable!(),
                };
                Ok(())
            },
        );
        assert!(match_res.is_ok());
        assert!(unmount);
    }

    #[test]
    fn name_value_pair() {
        let arguments = [Argument::short_value(
            'c',
            "cpus",
            "N",
            "Number of CPUs to use. (default: 1)",
        )];
        let match_res = set_arguments(
            ["-c", "5", "--cpus", "5", "-c5", "--cpus=5"].iter(),
            &arguments[..],
            |name, value| {
                assert_eq!(name, "cpus");
                assert_eq!(value, Some("5"));
                Ok(())
            },
        );
        assert!(match_res.is_ok());
    }

    #[test]
    fn flag_or_value() {
        let run_case = |args| -> Option<String> {
            let arguments = [
                Argument::positional("FILES", "files to operate on"),
                Argument::flag_or_value("gpu", "[2D|3D]", "Enable or configure gpu"),
                Argument::flag("foo", "Enable foo."),
                Argument::value("bar", "stuff", "Configure bar."),
            ];

            let mut gpu_value: Option<String> = None;
            let match_res =
                set_arguments(args, &arguments[..], |name: &str, value: Option<&str>| {
                    match name {
                        "" => assert_eq!(value.unwrap(), "file1"),
                        "foo" => assert!(value.is_none()),
                        "bar" => assert_eq!(value.unwrap(), "stuff"),
                        "gpu" => match value {
                            Some(v) => {
                                match v {
                                    "2D" => {
                                        gpu_value = Some("2D".to_string());
                                    }
                                    "3D" => {
                                        gpu_value = Some("3D".to_string());
                                    }
                                    _ => {
                                        return Err(Error::InvalidValue {
                                            value: v.to_string(),
                                            expected: "2D or 3D",
                                        })
                                    }
                                }
                                gpu_value = Some(v.to_string());
                            }
                            None => {
                                gpu_value = None;
                            }
                        },
                        _ => unreachable!(),
                    };
                    Ok(())
                });

            assert!(match_res.is_ok());
            gpu_value
        };

        // Used as flag and followed by positional
        assert_eq!(run_case(["--gpu", "file1"].iter()), None);
        // Used as flag and followed by flag
        assert_eq!(run_case(["--gpu", "--foo", "file1",].iter()), None);
        // Used as flag and followed by value
        assert_eq!(run_case(["--gpu", "--bar=stuff", "file1"].iter()), None);

        // Used as value and followed by positional
        assert_eq!(run_case(["--gpu=2D", "file1"].iter()).unwrap(), "2D");
        // Used as value and followed by flag
        assert_eq!(run_case(["--gpu=2D", "--foo"].iter()).unwrap(), "2D");
        // Used as value and followed by value
        assert_eq!(
            run_case(["--gpu=2D", "--bar=stuff", "file1"].iter()).unwrap(),
            "2D"
        );
    }
}
<|endoftext|>