<commit_before>use std::env;
use std::marker::PhantomData;
use clap::{self, Arg, SubCommand};
use shlex;

use app::{ClapApp, ClapRun};
use core::{Query, Repository, Workspace};
use util;
use vcs::{self, Vcs};

/// Toplevel application
pub enum Command<'a> {
  New(NewCommand<'a>),
  Add(AddCommand<'a>),
  Clone(CloneCommand<'a>),
  Scan(ScanCommand<'a>),
  List(ListCommand<'a>),
  Foreach(ForeachCommand<'a>),
}

impl<'a> ClapApp for Command<'a> {
  fn make_app<'b, 'c: 'b>(app: clap::App<'b, 'c>) -> clap::App<'b, 'c> {
    app.subcommand(NewCommand::make_app(SubCommand::with_name("new")))
       .subcommand(AddCommand::make_app(SubCommand::with_name("add")))
       .subcommand(CloneCommand::make_app(SubCommand::with_name("clone")))
       .subcommand(ListCommand::make_app(SubCommand::with_name("list")))
       .subcommand(ForeachCommand::make_app(SubCommand::with_name("foreach")))
       .subcommand(ScanCommand::make_app(SubCommand::with_name("scan")))
  }
}

impl<'a, 'b: 'a> From<&'b clap::ArgMatches<'a>> for Command<'a> {
  fn from(m: &'b clap::ArgMatches<'a>) -> Command<'a> {
    match m.subcommand() {
      ("new", Some(m)) => Command::New(m.into()),
      ("add", Some(m)) => Command::Add(m.into()),
      ("clone", Some(m)) => Command::Clone(m.into()),
      ("list", Some(m)) => Command::List(m.into()),
      ("foreach", Some(m)) => Command::Foreach(m.into()),
      ("scan", Some(m)) => Command::Scan(m.into()),
      _ => unreachable!(),
    }
  }
}

impl<'a> Command<'a> {
  pub fn run(self) -> ::Result<()> {
    match self {
      Command::New(m) => m.run(),
      Command::Add(m) => m.run(),
      Command::Clone(m) => m.run(),
      Command::List(m) => m.run(),
      Command::Foreach(m) => m.run(),
      Command::Scan(m) => m.run(),
    }
  }
}

/// Subcommand `new`
pub struct NewCommand<'a> {
  query: &'a str,
  root: Option<&'a str>,
  dry_run: bool,
  vcs: Option<Vcs>,
}

impl<'a> ClapApp for NewCommand<'a> {
  fn make_app<'b, 'c: 'b>(app: clap::App<'b, 'c>) -> clap::App<'b, 'c> {
    app.about("Create a new Git repository with intuitive directory structure")
       .arg_from_usage("<query>          'URL or query of remote repository'")
       .arg_from_usage("--root=[root]    'Target root directory of repository")
       .arg_from_usage("-n, --dry-run    'Do not actually create a new repository'")
       .arg(Arg::from_usage("--vcs=[vcs] 'Used Version Control System'").possible_values(Vcs::possible_values()))
  }
}

impl<'a, 'b: 'a> From<&'b clap::ArgMatches<'a>> for NewCommand<'a> {
  fn from(m: &'b clap::ArgMatches<'a>) -> NewCommand<'a> {
    NewCommand {
      query: m.value_of("query").unwrap(),
      root: m.value_of("root"),
      dry_run: m.is_present("dry-run"),
      vcs: m.value_of("vcs").and_then(|s| s.parse().ok()),
    }
  }
}

impl<'a> ClapRun for NewCommand<'a> {
  fn run(self) -> ::Result<()> {
    let mut workspace = Workspace::new(self.root)?;
    let root = workspace.root_dir().ok_or("Unknown root directory")?;

    let query: Query = self.query.parse()?;
    let path = query.to_local_path()?;
    let path = root.join(path);
    if vcs::detect_from_path(&path).is_some() {
      println!("The repository {} has already existed.", path.display());
      return Ok(());
    }

    let vcs = self.vcs.unwrap_or(Vcs::Git);

    print!("Creating an empty repository at \"{}\"", path.display());
    print!(" (VCS: {:?})", vcs);
    println!();

    if !self.dry_run {
      vcs.do_init(&path)?;
      let repo = Repository::from_path(path)?;
      workspace.add_repository(repo);
      workspace.save_cache()?;
    }

    Ok(())
  }
}

pub struct AddCommand<'a> {
  path: Option<&'a str>,
}

impl<'a> ClapApp for AddCommand<'a> {
  fn make_app<'b, 'c: 'b>(app: clap::App<'b, 'c>) -> clap::App<'b, 'c> {
    app.about("Add existed repository under management")
       .arg_from_usage("[path]  'Path of local repository'")
  }
}

impl<'a, 'b: 'a> From<&'b clap::ArgMatches<'a>> for AddCommand<'a> {
  fn from(m: &'b clap::ArgMatches<'a>) -> AddCommand<'a> {
    AddCommand { path: m.value_of("path") }
  }
}

impl<'a> ClapRun for AddCommand<'a> {
  fn run(self) -> ::Result<()> {
    let mut workspace = Workspace::new(None)?;

    let mut path = if let Some(path) = self.path {
      util::make_path_buf(path)?
    } else {
      env::current_dir()?
    };
    if !path.is_absolute() {
      path = env::current_dir()?.join(path);
    }

    let repo = Repository::from_path(path)?;
    if !repo.is_vcs() {
      Err("Given path is not a repository")?;
    }
    workspace.add_repository(repo);
    workspace.save_cache()?;

    Ok(())
  }
}

/// Subcommand `clone`
pub struct CloneCommand<'a> {
  query: Option<&'a str>,
  arg: Option<&'a str>,
  root: Option<&'a str>,
  dry_run: bool,
  ssh: bool,
  vcs: Option<Vcs>,
}

impl<'a> ClapApp for CloneCommand<'a> {
  fn make_app<'b, 'c: 'b>(app: clap::App<'b, 'c>) -> clap::App<'b, 'c> {
    app.about("Clone remote repositories into the root directory")
       .arg_from_usage("[query]         'URL or query of remote repository'")
       .arg_from_usage("--root=[root]   'Target root directory of cloned repository'")
       .arg_from_usage("--arg=[arg]     'Supplemental arguments for Git command'")
       .arg_from_usage("-n, --dry-run   'Do not actually execute Git command'")
       .arg_from_usage("-s, --ssh       'Use SSH protocol'")
       .arg(Arg::from_usage("--vcs=[vcs] 'Used Version Control System'").possible_values(Vcs::possible_values()))
  }
}

impl<'a, 'b: 'a> From<&'b clap::ArgMatches<'a>> for CloneCommand<'a> {
  fn from(m: &'b clap::ArgMatches<'a>) -> CloneCommand<'a> {
    CloneCommand {
      query: m.value_of("query"),
      arg: m.value_of("arg"),
      root: m.value_of("root"),
      dry_run: m.is_present("dry-run"),
      ssh: m.is_present("ssh"),
      vcs: m.value_of("vcs").and_then(|s| s.parse().ok()),
    }
  }
}

impl<'a> ClapRun for CloneCommand<'a> {
  fn run(self) -> ::Result<()> {
    let args = self.arg
                   .and_then(|s| shlex::split(s))
                   .unwrap_or_default();

    let mut workspace = Workspace::new(self.root)?;
    let root = workspace.root_dir().ok_or("Unknown root directory")?;

    for query in self.collect_queries()? {
      let path = query.to_local_path()?;
      let path = root.join(path);
      if vcs::detect_from_path(&path).is_some() {
        println!("The repository {} has already existed.", path.display());
        continue;
      }

      let url = query.to_url(self.ssh)?;

      let vcs = self.vcs.unwrap_or(Vcs::Git);

      println!("Clone: \"{}\" \"{}\" {}, {:?}",
               url,
               path.display(),
               util::join_str(&args),
               vcs);

      if !self.dry_run {
        vcs.do_clone(&path, &url, &args)?;
        let mut repo = Repository::from_path(path)?;
        repo.set_url(url);

        workspace.add_repository(repo);
      }
    }

    if !self.dry_run {
      workspace.save_cache()?;
    }

    Ok(())
  }
}

impl<'a> CloneCommand<'a> {
  fn collect_queries(&self) -> ::Result<Vec<Query>> {
    let mut queries = Vec::new();

    if let Some(query) = self.query {
      let query = query.parse()?;
      queries.push(query);

    } else {
      use std::io::BufRead;

      let stdin = ::std::io::stdin();
      for query in stdin.lock().lines().filter_map(|l| l.ok()) {
        let query = query.parse()?;
        queries.push(query);
      }
    }

    Ok(queries)
  }
}

/// Subcommand `scan`
pub struct ScanCommand<'a> {
  verbose: bool,
  prune: bool,
  depth: Option<usize>,
  marker: PhantomData<&'a usize>,
}

impl<'a> ClapApp for ScanCommand<'a> {
  fn make_app<'b, 'c: 'b>(app: clap::App<'b, 'c>) -> clap::App<'b, 'c> {
    app.about("Scan directories to create cache of repositories list")
       .arg_from_usage("-v, --verbose    'Use verbose output'")
       .arg_from_usage("-p, --prune      'Ignore repositories located at outside of base directories'")
       .arg_from_usage("--depth=[depth]  'Maximal depth of entries for each base directory'")
  }
}

impl<'a, 'b: 'a> From<&'b clap::ArgMatches<'a>> for ScanCommand<'a> {
  fn from(m: &'b clap::ArgMatches<'a>) -> ScanCommand<'a> {
    ScanCommand {
      verbose: m.is_present("verbose"),
      prune: m.is_present("prune"),
      depth: m.value_of("depth").and_then(|s| s.parse().ok()),
      marker: PhantomData,
    }
  }
}

impl<'a> ClapRun for ScanCommand<'a> {
  fn run(self) -> ::Result<()> {
    let mut workspace = Workspace::new(None)?;
    workspace.scan_repositories(self.verbose, self.prune, self.depth)?;
    workspace.save_cache()?;
    Ok(())
  }
}

/// Subcommand `list`
pub struct ListCommand<'a> {
  marker: PhantomData<&'a usize>,
}

impl<'a> ClapApp for ListCommand<'a> {
  fn make_app<'b, 'c: 'b>(app: clap::App<'b, 'c>) -> clap::App<'b, 'c> {
    app.about("List local repositories managed by rhq")
  }
}

impl<'a, 'b: 'a> From<&'b clap::ArgMatches<'a>> for ListCommand<'a> {
  fn from(_: &'b clap::ArgMatches<'a>) -> ListCommand<'a> {
    ListCommand { marker: PhantomData }
  }
}

impl<'a> ClapRun for ListCommand<'a> {
  fn run(self) -> ::Result<()> {
    let workspace = Workspace::new(None)?;
    let repos = workspace.repositories()
                         .ok_or("The cache has not initialized yet")?;
    for repo in repos {
      println!("{}", repo.path_string());
    }
    Ok(())
  }
}

/// Subcommand `foreach`
pub struct ForeachCommand<'a> {
  command: &'a str,
  args: Option<clap::Values<'a>>,
  dry_run: bool,
}

impl<'a> ClapApp for ForeachCommand<'a> {
  fn make_app<'b, 'c: 'b>(app: clap::App<'b, 'c>) -> clap::App<'b, 'c> {
    app.about("Execute command into each repositories")
       .arg_from_usage("<command>       'Command name'")
       .arg_from_usage("[args]...       'Arguments of command'")
       .arg_from_usage("-n, --dry-run   'Do not actually execute command'")
  }
}

impl<'a, 'b: 'a> From<&'b clap::ArgMatches<'a>> for ForeachCommand<'a> {
  fn from(m: &'b clap::ArgMatches<'a>) -> ForeachCommand<'a> {
    ForeachCommand {
      command: m.value_of("command").unwrap(),
      args: m.values_of("args"),
      dry_run: m.is_present("dry-run"),
    }
  }
}

impl<'a> ClapRun for ForeachCommand<'a> {
  fn run(self) -> ::Result<()> {
    let args: Vec<_> = self.args.map(|s| s.collect()).unwrap_or_default();
    let workspace = Workspace::new(None)?;
    let repos = workspace.repositories()
                         .ok_or("The cache has not initialized yet")?;
    for repo in repos {
      if self.dry_run {
        println!("+ {} {}", self.command, util::join_str(&args));
      } else {
        repo.run_command(self.command, &args)?;
      }
    }
    Ok(())
  }
}
<commit_msg>rhq clone: make <query> required<commit_after>use std::env;
use std::marker::PhantomData;
use clap::{self, Arg, SubCommand};
use shlex;

use app::{ClapApp, ClapRun};
use core::{Query, Repository, Workspace};
use util;
use vcs::{self, Vcs};

/// Toplevel application
pub enum Command<'a> {
  New(NewCommand<'a>),
  Add(AddCommand<'a>),
  Clone(CloneCommand<'a>),
  Scan(ScanCommand<'a>),
  List(ListCommand<'a>),
  Foreach(ForeachCommand<'a>),
}

impl<'a> ClapApp for Command<'a> {
  fn make_app<'b, 'c: 'b>(app: clap::App<'b, 'c>) -> clap::App<'b, 'c> {
    app.subcommand(NewCommand::make_app(SubCommand::with_name("new")))
       .subcommand(AddCommand::make_app(SubCommand::with_name("add")))
       .subcommand(CloneCommand::make_app(SubCommand::with_name("clone")))
       .subcommand(ListCommand::make_app(SubCommand::with_name("list")))
       .subcommand(ForeachCommand::make_app(SubCommand::with_name("foreach")))
       .subcommand(ScanCommand::make_app(SubCommand::with_name("scan")))
  }
}

impl<'a, 'b: 'a> From<&'b clap::ArgMatches<'a>> for Command<'a> {
  fn from(m: &'b clap::ArgMatches<'a>) -> Command<'a> {
    match m.subcommand() {
      ("new", Some(m)) => Command::New(m.into()),
      ("add", Some(m)) => Command::Add(m.into()),
      ("clone", Some(m)) => Command::Clone(m.into()),
      ("list", Some(m)) => Command::List(m.into()),
      ("foreach", Some(m)) => Command::Foreach(m.into()),
      ("scan", Some(m)) => Command::Scan(m.into()),
      _ => unreachable!(),
    }
  }
}

impl<'a> Command<'a> {
  pub fn run(self) -> ::Result<()> {
    match self {
      Command::New(m) => m.run(),
      Command::Add(m) => m.run(),
      Command::Clone(m) => m.run(),
      Command::List(m) => m.run(),
      Command::Foreach(m) => m.run(),
      Command::Scan(m) => m.run(),
    }
  }
}

/// Subcommand `new`
pub struct NewCommand<'a> {
  query: &'a str,
  root: Option<&'a str>,
  dry_run: bool,
  vcs: Option<Vcs>,
}

impl<'a> ClapApp for NewCommand<'a> {
  fn make_app<'b, 'c: 'b>(app: clap::App<'b, 'c>) -> clap::App<'b, 'c> {
    app.about("Create a new Git repository with intuitive directory structure")
       .arg_from_usage("<query>          'URL or query of remote repository'")
       .arg_from_usage("--root=[root]    'Target root directory of repository")
       .arg_from_usage("-n, --dry-run    'Do not actually create a new repository'")
       .arg(Arg::from_usage("--vcs=[vcs] 'Used Version Control System'").possible_values(Vcs::possible_values()))
  }
}

impl<'a, 'b: 'a> From<&'b clap::ArgMatches<'a>> for NewCommand<'a> {
  fn from(m: &'b clap::ArgMatches<'a>) -> NewCommand<'a> {
    NewCommand {
      query: m.value_of("query").unwrap(),
      root: m.value_of("root"),
      dry_run: m.is_present("dry-run"),
      vcs: m.value_of("vcs").and_then(|s| s.parse().ok()),
    }
  }
}

impl<'a> ClapRun for NewCommand<'a> {
  fn run(self) -> ::Result<()> {
    let mut workspace = Workspace::new(self.root)?;
    let root = workspace.root_dir().ok_or("Unknown root directory")?;

    let query: Query = self.query.parse()?;
    let path = query.to_local_path()?;
    let path = root.join(path);
    if vcs::detect_from_path(&path).is_some() {
      println!("The repository {} has already existed.", path.display());
      return Ok(());
    }

    let vcs = self.vcs.unwrap_or(Vcs::Git);

    print!("Creating an empty repository at \"{}\"", path.display());
    print!(" (VCS: {:?})", vcs);
    println!();

    if !self.dry_run {
      vcs.do_init(&path)?;
      let repo = Repository::from_path(path)?;
      workspace.add_repository(repo);
      workspace.save_cache()?;
    }

    Ok(())
  }
}

pub struct AddCommand<'a> {
  path: Option<&'a str>,
}

impl<'a> ClapApp for AddCommand<'a> {
  fn make_app<'b, 'c: 'b>(app: clap::App<'b, 'c>) -> clap::App<'b, 'c> {
    app.about("Add existed repository under management")
       .arg_from_usage("[path]  'Path of local repository'")
  }
}

impl<'a, 'b: 'a> From<&'b clap::ArgMatches<'a>> for AddCommand<'a> {
  fn from(m: &'b clap::ArgMatches<'a>) -> AddCommand<'a> {
    AddCommand { path: m.value_of("path") }
  }
}

impl<'a> ClapRun for AddCommand<'a> {
  fn run(self) -> ::Result<()> {
    let mut workspace = Workspace::new(None)?;

    let mut path = if let Some(path) = self.path {
      util::make_path_buf(path)?
    } else {
      env::current_dir()?
    };
    if !path.is_absolute() {
      path = env::current_dir()?.join(path);
    }

    let repo = Repository::from_path(path)?;
    if !repo.is_vcs() {
      Err("Given path is not a repository")?;
    }
    workspace.add_repository(repo);
    workspace.save_cache()?;

    Ok(())
  }
}

/// Subcommand `clone`
pub struct CloneCommand<'a> {
  query: Query,
  arg: Option<&'a str>,
  root: Option<&'a str>,
  dry_run: bool,
  ssh: bool,
  vcs: Option<Vcs>,
}

impl<'a> ClapApp for CloneCommand<'a> {
  fn make_app<'b, 'c: 'b>(app: clap::App<'b, 'c>) -> clap::App<'b, 'c> {
    app.about("Clone remote repositories into the root directory")
       .arg_from_usage("<query>         'URL or query of remote repository'")
       .arg_from_usage("--root=[root]   'Target root directory of cloned repository'")
       .arg_from_usage("--arg=[arg]     'Supplemental arguments for Git command'")
       .arg_from_usage("-n, --dry-run   'Do not actually execute Git command'")
       .arg_from_usage("-s, --ssh       'Use SSH protocol'")
       .arg(Arg::from_usage("--vcs=[vcs] 'Used Version Control System'").possible_values(Vcs::possible_values()))
  }
}

impl<'a, 'b: 'a> From<&'b clap::ArgMatches<'a>> for CloneCommand<'a> {
  fn from(m: &'b clap::ArgMatches<'a>) -> CloneCommand<'a> {
    CloneCommand {
      query: m.value_of("query")
              .and_then(|s| s.parse().ok())
              .unwrap(),
      arg: m.value_of("arg"),
      root: m.value_of("root"),
      dry_run: m.is_present("dry-run"),
      ssh: m.is_present("ssh"),
      vcs: m.value_of("vcs").and_then(|s| s.parse().ok()),
    }
  }
}

impl<'a> ClapRun for CloneCommand<'a> {
  fn run(self) -> ::Result<()> {
    let args = self.arg
                   .and_then(|s| shlex::split(s))
                   .unwrap_or_default();

    let mut workspace = Workspace::new(self.root)?;
    let root = workspace.root_dir().ok_or("Unknown root directory")?;

    let path = self.query.to_local_path()?;
    let path = root.join(path);
    if vcs::detect_from_path(&path).is_some() {
      println!("The repository {} has already existed.", path.display());
      return Ok(());
    }

    let url = self.query.to_url(self.ssh)?;

    let vcs = self.vcs.unwrap_or(Vcs::Git);

    println!("Clone: \"{}\" \"{}\" {}, {:?}",
             url,
             path.display(),
             util::join_str(&args),
             vcs);

    if !self.dry_run {
      vcs.do_clone(&path, &url, &args)?;

      let mut repo = Repository::from_path(path)?;
      repo.set_url(url);
      workspace.add_repository(repo);
      workspace.save_cache()?;
    }

    Ok(())
  }
}

/// Subcommand `scan`
pub struct ScanCommand<'a> {
  verbose: bool,
  prune: bool,
  depth: Option<usize>,
  marker: PhantomData<&'a usize>,
}

impl<'a> ClapApp for ScanCommand<'a> {
  fn make_app<'b, 'c: 'b>(app: clap::App<'b, 'c>) -> clap::App<'b, 'c> {
    app.about("Scan directories to create cache of repositories list")
       .arg_from_usage("-v, --verbose    'Use verbose output'")
       .arg_from_usage("-p, --prune      'Ignore repositories located at outside of base directories'")
       .arg_from_usage("--depth=[depth]  'Maximal depth of entries for each base directory'")
  }
}

impl<'a, 'b: 'a> From<&'b clap::ArgMatches<'a>> for ScanCommand<'a> {
  fn from(m: &'b clap::ArgMatches<'a>) -> ScanCommand<'a> {
    ScanCommand {
      verbose: m.is_present("verbose"),
      prune: m.is_present("prune"),
      depth: m.value_of("depth").and_then(|s| s.parse().ok()),
      marker: PhantomData,
    }
  }
}

impl<'a> ClapRun for ScanCommand<'a> {
  fn run(self) -> ::Result<()> {
    let mut workspace = Workspace::new(None)?;
    workspace.scan_repositories(self.verbose, self.prune, self.depth)?;
    workspace.save_cache()?;
    Ok(())
  }
}

/// Subcommand `list`
pub struct ListCommand<'a> {
  marker: PhantomData<&'a usize>,
}

impl<'a> ClapApp for ListCommand<'a> {
  fn make_app<'b, 'c: 'b>(app: clap::App<'b, 'c>) -> clap::App<'b, 'c> {
    app.about("List local repositories managed by rhq")
  }
}

impl<'a, 'b: 'a> From<&'b clap::ArgMatches<'a>> for ListCommand<'a> {
  fn from(_: &'b clap::ArgMatches<'a>) -> ListCommand<'a> {
    ListCommand { marker: PhantomData }
  }
}

impl<'a> ClapRun for ListCommand<'a> {
  fn run(self) -> ::Result<()> {
    let workspace = Workspace::new(None)?;
    let repos = workspace.repositories()
                         .ok_or("The cache has not initialized yet")?;
    for repo in repos {
      println!("{}", repo.path_string());
    }
    Ok(())
  }
}

/// Subcommand `foreach`
pub struct ForeachCommand<'a> {
  command: &'a str,
  args: Option<clap::Values<'a>>,
  dry_run: bool,
}

impl<'a> ClapApp for ForeachCommand<'a> {
  fn make_app<'b, 'c: 'b>(app: clap::App<'b, 'c>) -> clap::App<'b, 'c> {
    app.about("Execute command into each repositories")
       .arg_from_usage("<command>       'Command name'")
       .arg_from_usage("[args]...       'Arguments of command'")
       .arg_from_usage("-n, --dry-run   'Do not actually execute command'")
  }
}

impl<'a, 'b: 'a> From<&'b clap::ArgMatches<'a>> for ForeachCommand<'a> {
  fn from(m: &'b clap::ArgMatches<'a>) -> ForeachCommand<'a> {
    ForeachCommand {
      command: m.value_of("command").unwrap(),
      args: m.values_of("args"),
      dry_run: m.is_present("dry-run"),
    }
  }
}

impl<'a> ClapRun for ForeachCommand<'a> {
  fn run(self) -> ::Result<()> {
    let args: Vec<_> = self.args.map(|s| s.collect()).unwrap_or_default();
    let workspace = Workspace::new(None)?;
    let repos = workspace.repositories()
                         .ok_or("The cache has not initialized yet")?;
    for repo in repos {
      if self.dry_run {
        println!("+ {} {}", self.command, util::join_str(&args));
      } else {
        repo.run_command(self.command, &args)?;
      }
    }
    Ok(())
  }
}
<|endoftext|>


<commit_before>use dns_parser::{self, QueryClass, QueryType, Name, RRData};
use std::collections::VecDeque;
use std::io;
use std::io::ErrorKind::WouldBlock;
use std::marker::PhantomData;
use std::net::{IpAddr, SocketAddr};
use futures::{Poll, Async, Future, Stream};
use futures::sync::mpsc;
use tokio::net::UdpSocket;
use tokio::reactor::Handle;

use super::{DEFAULT_TTL, MDNS_PORT};
use address_family::AddressFamily;
use net;
use services::{Services, ServiceData};

pub type AnswerBuilder = dns_parser::Builder<dns_parser::Answers>;

#[derive(Clone, Debug)]
pub enum Command {
    SendUnsolicited {
        svc: ServiceData,
        ttl: u32,
        include_ip: bool
    },
    Shutdown,
}

pub struct FSM<AF: AddressFamily> {
    socket: UdpSocket,
    services: Services,
    commands: mpsc::UnboundedReceiver<Command>,
    outgoing: VecDeque<(Vec<u8>, SocketAddr)>,
    _af: PhantomData<AF>,
}

impl <AF: AddressFamily> FSM<AF> {
    pub fn new(handle: &Handle, services: &Services)
        -> io::Result<(FSM<AF>, mpsc::UnboundedSender<Command>)>
    {
        let std_socket = AF::bind()?;
        let socket = UdpSocket::from_socket(std_socket, handle)?;
        let (tx, rx) = mpsc::unbounded();

        let fsm = FSM {
            socket: socket,
            services: services.clone(),
            commands: rx,
            outgoing: VecDeque::new(),
            _af: PhantomData,
        };

        Ok((fsm, tx))
    }

    fn recv_packets(&mut self) -> io::Result<()> {
        let mut buf = [0u8; 4096];
        loop {
            let (bytes, addr) = match self.socket.recv_from(&mut buf) {
                Ok((bytes, addr)) => (bytes, addr),
                Err(ref ioerr) if ioerr.kind() == WouldBlock => break,
                Err(err) => return Err(err),
            };

            if bytes >= buf.len() {
                warn!("buffer too small for packet from {:?}", addr);
                continue;
            }

            self.handle_packet(&buf[..bytes], addr);
        }
        return Ok(())
    }

    fn handle_packet(&mut self, buffer: &[u8], addr: SocketAddr) {
        trace!("received packet from {:?}", addr);

        let packet = match dns_parser::Packet::parse(buffer) {
            Ok(packet) => packet,
            Err(error) => {
                warn!("couldn't parse packet from {:?}: {}", addr, error);
                return;
            }
        };

        if !packet.header.query {
            trace!("received packet from {:?} with no query", addr);
            return;
        }

        if packet.header.truncated {
            warn!("dropping truncated packet from {:?}", addr);
            return;
        }

        let mut unicast_builder = dns_parser::Builder::new_response(packet.header.id, false).move_to::<dns_parser::Answers>();
        let mut multicast_builder = dns_parser::Builder::new_response(packet.header.id, false).move_to::<dns_parser::Answers>();
        unicast_builder.set_max_size(None);
        multicast_builder.set_max_size(None);

        for question in packet.questions {
            if question.qclass == QueryClass::IN || question.qclass == QueryClass::Any {
                if question.qu {
                    unicast_builder = self.handle_question(&question, unicast_builder);
                } else {
                    multicast_builder = self.handle_question(&question, multicast_builder);
                }
            }
        }

        if !multicast_builder.is_empty() {
            let response = multicast_builder.build().unwrap_or_else(|x| x);
            let addr = SocketAddr::new(AF::mdns_group(), MDNS_PORT);
            self.outgoing.push_back((response, addr));
        }

        if !unicast_builder.is_empty() {
            let response = unicast_builder.build().unwrap_or_else(|x| x);
            self.outgoing.push_back((response, addr));
        }
    }

    fn handle_question(&self, question: &dns_parser::Question, mut builder: AnswerBuilder) -> AnswerBuilder {
        let services = self.services.read().unwrap();

        match question.qtype {
            QueryType::A |
                QueryType::AAAA |
                QueryType::All if question.qname == *services.get_hostname() => {
                    builder = self.add_ip_rr(services.get_hostname(), builder, DEFAULT_TTL);
                }
            QueryType::PTR => {
                for svc in services.find_by_type(&question.qname) {
                    builder = svc.add_ptr_rr(builder, DEFAULT_TTL);
                    builder = svc.add_srv_rr(services.get_hostname(), builder, DEFAULT_TTL);
                    builder = svc.add_txt_rr(builder, DEFAULT_TTL);
                    builder = self.add_ip_rr(services.get_hostname(), builder, DEFAULT_TTL);
                }
            }
            QueryType::SRV => {
                if let Some(svc) = services.find_by_name(&question.qname) {
                    builder = svc.add_srv_rr(services.get_hostname(), builder, DEFAULT_TTL);
                    builder = self.add_ip_rr(services.get_hostname(), builder, DEFAULT_TTL);
                }
            }
            QueryType::TXT => {
                if let Some(svc) = services.find_by_name(&question.qname) {
                    builder = svc.add_txt_rr(builder, DEFAULT_TTL);
                }
            }
            _ => ()
        }

        builder
    }

    fn add_ip_rr(&self, hostname: &Name, mut builder: AnswerBuilder, ttl: u32) -> AnswerBuilder {
        for iface in net::getifaddrs() {
            if iface.is_loopback() {
                continue;
            }

            match iface.ip() {
                Some(IpAddr::V4(ip)) if !AF::v6() => {
                    builder = builder.add_answer(hostname, QueryClass::IN, ttl, &RRData::A(ip))
                }
                Some(IpAddr::V6(ip)) if AF::v6() => {
                    builder = builder.add_answer(hostname, QueryClass::IN, ttl, &RRData::AAAA(ip))
                }
                _ => ()
            }
        }

        builder
    }

    fn send_unsolicited(&mut self, svc: &ServiceData, ttl: u32, include_ip: bool) {
        let mut builder = dns_parser::Builder::new_response(0, false).move_to::<dns_parser::Answers>();
        builder.set_max_size(None);

        let services = self.services.read().unwrap();

        builder = svc.add_ptr_rr(builder, ttl);
        builder = svc.add_srv_rr(services.get_hostname(), builder, ttl);
        builder = svc.add_txt_rr(builder, ttl);
        if include_ip {
            builder = self.add_ip_rr(services.get_hostname(), builder, ttl);
        }

        if !builder.is_empty() {
            let response = builder.build().unwrap_or_else(|x| x);
            let addr = SocketAddr::new(AF::mdns_group(), MDNS_PORT);
            self.outgoing.push_back((response, addr));
        }
    }
}

/*
   impl rotor::Machine for FSM {
    type Context = ();
    type Seed = void::Void;

    fn create(seed: Self::Seed, _scope: &mut Scope<Self::Context>) -> rotor::Response<Self, rotor::Void> {
        void::unreachable(seed)
    }

    fn ready(self, _events: EventSet, _scope: &mut Scope<Self::Context>) -> rotor::Response<Self, Self::Seed> {
        self.recv_packets().unwrap();
        rotor::Response::ok(self)
    }

    fn spawned(self, _scope: &mut Scope<Self::Context>) -> rotor::Response<Self, Self::Seed> {
        unimplemented!()
    }

    fn timeout(self, _scope: &mut Scope<Self::Context>) -> rotor::Response<Self, Self::Seed> {
        unimplemented!()
    }

    fn wakeup(self, scope: &mut Scope<Self::Context>) -> rotor::Response<Self, Self::Seed> {
        loop {
            match self.rx.try_recv() {
                Ok(Command::Shutdown) => {
                    scope.shutdown_loop();
                    return rotor::Response::done();
                }
                Ok(Command::SendUnsolicited { svc, ttl, include_ip }) => {
                    match self.send_unsolicited(&svc, ttl, include_ip) {
                        Ok(_) => (),
                        Err(e) => {
                            warn!("Error sending unsolicited: {:?}", e);
                            return rotor::Response::error(Box::new(e));
                        }
                    }
                }
                Err(TryRecvError::Disconnected) => {
                    warn!("responder disconnected without shutdown");
                    scope.shutdown_loop();
                    return rotor::Response::done();
                }
                Err(TryRecvError::Empty) => {
                    break;
                }
            }
        }

        rotor::Response::ok(self)
    }
}
*/

impl <AF: AddressFamily> Future for FSM<AF> {
    type Item = ();
    type Error = io::Error;
    fn poll(&mut self) -> Poll<(), io::Error> {
        while let Async::Ready(cmd) = self.commands.poll().unwrap() {
            println!("cmd={:?}", cmd);
            match cmd {
                Some(Command::Shutdown) => return Ok(Async::Ready(())),
                Some(Command::SendUnsolicited { svc, ttl, include_ip }) => {
                    self.send_unsolicited(&svc, ttl, include_ip);
                }
                None => {
                    warn!("responder disconnected without shutdown");
                    return Ok(Async::Ready(()));
                }
            }
        }

        while let Async::Ready(()) = self.socket.poll_read() {
            self.recv_packets()?;
        }

        loop {
            if let Some(&(ref response, ref addr)) = self.outgoing.front() {
                trace!("sending packet to {:?}", addr);

                match self.socket.send_to(response, addr) {
                    Ok(_) => (),
                    Err(ref ioerr) if ioerr.kind() == WouldBlock => break,
                    Err(err) => warn!("error sending packet {:?}", err),
                }
            } else {
                break;
            }

            self.outgoing.pop_front();
        }

        Ok(Async::NotReady)
    }
}
<commit_msg>Remove leftover debug message<commit_after>use dns_parser::{self, QueryClass, QueryType, Name, RRData};
use std::collections::VecDeque;
use std::io;
use std::io::ErrorKind::WouldBlock;
use std::marker::PhantomData;
use std::net::{IpAddr, SocketAddr};
use futures::{Poll, Async, Future, Stream};
use futures::sync::mpsc;
use tokio::net::UdpSocket;
use tokio::reactor::Handle;

use super::{DEFAULT_TTL, MDNS_PORT};
use address_family::AddressFamily;
use net;
use services::{Services, ServiceData};

pub type AnswerBuilder = dns_parser::Builder<dns_parser::Answers>;

#[derive(Clone, Debug)]
pub enum Command {
    SendUnsolicited {
        svc: ServiceData,
        ttl: u32,
        include_ip: bool
    },
    Shutdown,
}

pub struct FSM<AF: AddressFamily> {
    socket: UdpSocket,
    services: Services,
    commands: mpsc::UnboundedReceiver<Command>,
    outgoing: VecDeque<(Vec<u8>, SocketAddr)>,
    _af: PhantomData<AF>,
}

impl <AF: AddressFamily> FSM<AF> {
    pub fn new(handle: &Handle, services: &Services)
        -> io::Result<(FSM<AF>, mpsc::UnboundedSender<Command>)>
    {
        let std_socket = AF::bind()?;
        let socket = UdpSocket::from_socket(std_socket, handle)?;
        let (tx, rx) = mpsc::unbounded();

        let fsm = FSM {
            socket: socket,
            services: services.clone(),
            commands: rx,
            outgoing: VecDeque::new(),
            _af: PhantomData,
        };

        Ok((fsm, tx))
    }

    fn recv_packets(&mut self) -> io::Result<()> {
        let mut buf = [0u8; 4096];
        loop {
            let (bytes, addr) = match self.socket.recv_from(&mut buf) {
                Ok((bytes, addr)) => (bytes, addr),
                Err(ref ioerr) if ioerr.kind() == WouldBlock => break,
                Err(err) => return Err(err),
            };

            if bytes >= buf.len() {
                warn!("buffer too small for packet from {:?}", addr);
                continue;
            }

            self.handle_packet(&buf[..bytes], addr);
        }
        return Ok(())
    }

    fn handle_packet(&mut self, buffer: &[u8], addr: SocketAddr) {
        trace!("received packet from {:?}", addr);

        let packet = match dns_parser::Packet::parse(buffer) {
            Ok(packet) => packet,
            Err(error) => {
                warn!("couldn't parse packet from {:?}: {}", addr, error);
                return;
            }
        };

        if !packet.header.query {
            trace!("received packet from {:?} with no query", addr);
            return;
        }

        if packet.header.truncated {
            warn!("dropping truncated packet from {:?}", addr);
            return;
        }

        let mut unicast_builder = dns_parser::Builder::new_response(packet.header.id, false).move_to::<dns_parser::Answers>();
        let mut multicast_builder = dns_parser::Builder::new_response(packet.header.id, false).move_to::<dns_parser::Answers>();
        unicast_builder.set_max_size(None);
        multicast_builder.set_max_size(None);

        for question in packet.questions {
            if question.qclass == QueryClass::IN || question.qclass == QueryClass::Any {
                if question.qu {
                    unicast_builder = self.handle_question(&question, unicast_builder);
                } else {
                    multicast_builder = self.handle_question(&question, multicast_builder);
                }
            }
        }

        if !multicast_builder.is_empty() {
            let response = multicast_builder.build().unwrap_or_else(|x| x);
            let addr = SocketAddr::new(AF::mdns_group(), MDNS_PORT);
            self.outgoing.push_back((response, addr));
        }

        if !unicast_builder.is_empty() {
            let response = unicast_builder.build().unwrap_or_else(|x| x);
            self.outgoing.push_back((response, addr));
        }
    }

    fn handle_question(&self, question: &dns_parser::Question, mut builder: AnswerBuilder) -> AnswerBuilder {
        let services = self.services.read().unwrap();

        match question.qtype {
            QueryType::A |
                QueryType::AAAA |
                QueryType::All if question.qname == *services.get_hostname() => {
                    builder = self.add_ip_rr(services.get_hostname(), builder, DEFAULT_TTL);
                }
            QueryType::PTR => {
                for svc in services.find_by_type(&question.qname) {
                    builder = svc.add_ptr_rr(builder, DEFAULT_TTL);
                    builder = svc.add_srv_rr(services.get_hostname(), builder, DEFAULT_TTL);
                    builder = svc.add_txt_rr(builder, DEFAULT_TTL);
                    builder = self.add_ip_rr(services.get_hostname(), builder, DEFAULT_TTL);
                }
            }
            QueryType::SRV => {
                if let Some(svc) = services.find_by_name(&question.qname) {
                    builder = svc.add_srv_rr(services.get_hostname(), builder, DEFAULT_TTL);
                    builder = self.add_ip_rr(services.get_hostname(), builder, DEFAULT_TTL);
                }
            }
            QueryType::TXT => {
                if let Some(svc) = services.find_by_name(&question.qname) {
                    builder = svc.add_txt_rr(builder, DEFAULT_TTL);
                }
            }
            _ => ()
        }

        builder
    }

    fn add_ip_rr(&self, hostname: &Name, mut builder: AnswerBuilder, ttl: u32) -> AnswerBuilder {
        for iface in net::getifaddrs() {
            if iface.is_loopback() {
                continue;
            }

            match iface.ip() {
                Some(IpAddr::V4(ip)) if !AF::v6() => {
                    builder = builder.add_answer(hostname, QueryClass::IN, ttl, &RRData::A(ip))
                }
                Some(IpAddr::V6(ip)) if AF::v6() => {
                    builder = builder.add_answer(hostname, QueryClass::IN, ttl, &RRData::AAAA(ip))
                }
                _ => ()
            }
        }

        builder
    }

    fn send_unsolicited(&mut self, svc: &ServiceData, ttl: u32, include_ip: bool) {
        let mut builder = dns_parser::Builder::new_response(0, false).move_to::<dns_parser::Answers>();
        builder.set_max_size(None);

        let services = self.services.read().unwrap();

        builder = svc.add_ptr_rr(builder, ttl);
        builder = svc.add_srv_rr(services.get_hostname(), builder, ttl);
        builder = svc.add_txt_rr(builder, ttl);
        if include_ip {
            builder = self.add_ip_rr(services.get_hostname(), builder, ttl);
        }

        if !builder.is_empty() {
            let response = builder.build().unwrap_or_else(|x| x);
            let addr = SocketAddr::new(AF::mdns_group(), MDNS_PORT);
            self.outgoing.push_back((response, addr));
        }
    }
}

/*
   impl rotor::Machine for FSM {
    type Context = ();
    type Seed = void::Void;

    fn create(seed: Self::Seed, _scope: &mut Scope<Self::Context>) -> rotor::Response<Self, rotor::Void> {
        void::unreachable(seed)
    }

    fn ready(self, _events: EventSet, _scope: &mut Scope<Self::Context>) -> rotor::Response<Self, Self::Seed> {
        self.recv_packets().unwrap();
        rotor::Response::ok(self)
    }

    fn spawned(self, _scope: &mut Scope<Self::Context>) -> rotor::Response<Self, Self::Seed> {
        unimplemented!()
    }

    fn timeout(self, _scope: &mut Scope<Self::Context>) -> rotor::Response<Self, Self::Seed> {
        unimplemented!()
    }

    fn wakeup(self, scope: &mut Scope<Self::Context>) -> rotor::Response<Self, Self::Seed> {
        loop {
            match self.rx.try_recv() {
                Ok(Command::Shutdown) => {
                    scope.shutdown_loop();
                    return rotor::Response::done();
                }
                Ok(Command::SendUnsolicited { svc, ttl, include_ip }) => {
                    match self.send_unsolicited(&svc, ttl, include_ip) {
                        Ok(_) => (),
                        Err(e) => {
                            warn!("Error sending unsolicited: {:?}", e);
                            return rotor::Response::error(Box::new(e));
                        }
                    }
                }
                Err(TryRecvError::Disconnected) => {
                    warn!("responder disconnected without shutdown");
                    scope.shutdown_loop();
                    return rotor::Response::done();
                }
                Err(TryRecvError::Empty) => {
                    break;
                }
            }
        }

        rotor::Response::ok(self)
    }
}
*/

impl <AF: AddressFamily> Future for FSM<AF> {
    type Item = ();
    type Error = io::Error;
    fn poll(&mut self) -> Poll<(), io::Error> {
        while let Async::Ready(cmd) = self.commands.poll().unwrap() {
            match cmd {
                Some(Command::Shutdown) => return Ok(Async::Ready(())),
                Some(Command::SendUnsolicited { svc, ttl, include_ip }) => {
                    self.send_unsolicited(&svc, ttl, include_ip);
                }
                None => {
                    warn!("responder disconnected without shutdown");
                    return Ok(Async::Ready(()));
                }
            }
        }

        while let Async::Ready(()) = self.socket.poll_read() {
            self.recv_packets()?;
        }

        loop {
            if let Some(&(ref response, ref addr)) = self.outgoing.front() {
                trace!("sending packet to {:?}", addr);

                match self.socket.send_to(response, addr) {
                    Ok(_) => (),
                    Err(ref ioerr) if ioerr.kind() == WouldBlock => break,
                    Err(err) => warn!("error sending packet {:?}", err),
                }
            } else {
                break;
            }

            self.outgoing.pop_front();
        }

        Ok(Async::NotReady)
    }
}
<|endoftext|>


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

//! Capture the screen with xlib

#![allow(dead_code, non_upper_case_globals, non_camel_case_types)]

extern crate x11;
extern crate libc;

use libc::c_ulong;

pub mod ffi;

#[derive(Debug, Eq, PartialEq)]
struct RGB8 {
	r: u8,
	g: u8,
	b: u8,
}

/// From a bitmask for a color in a pixel, calculate the color size in bits and the bitshift
fn mask_size_and_shift(mut mask: c_ulong) -> (c_ulong, u16) {
	let (mut bits, mut shift) = (0, 0);
	let mut prev = 0;
	for _ in 0..(c_ulong::max_value() as f64).log2() as u16 {
		if mask & 1 == 1 {
			bits += 1;
			prev = 1;
		} else if prev == 1 {
			break;
		} else {
			shift += 1;
		}
		mask >>= 1;
	}
	(bits, shift)
}

#[test]
fn test_create_window() {
	use x11::xlib;
	use std::ptr;

	unsafe {

	// use the information from the environment variable DISPLAY to create the X connection:
	let display = match xlib::XOpenDisplay(ptr::null_mut()) {
		null if null.is_null() => panic!("Unable to connect X server"),
		display => display,
	};

	let screen_num = xlib::XDefaultScreen(display);
	let (black_pixel, white_pixel) = (xlib::XBlackPixel(display, screen_num),
		xlib::XWhitePixel(display, screen_num));

	// Once the display is initialized, create the window
	let (x, y) = (0, 0);
	let (width, height, border_width) = (400, 500, 5);
	let (border_color, background_color) = (black_pixel, white_pixel);
	let window = xlib::XCreateSimpleWindow(display,
		xlib::XRootWindow(display, screen_num),
		x, y,
		width, height, border_width,
		border_color, background_color);

	// this routine determines which types of input are allowed in the input.
	xlib::XSelectInput(display, window,
		xlib::ExposureMask | xlib::ButtonPressMask | xlib::KeyPressMask);

	// create the Graphics Context
	let graphics_context = xlib::XCreateGC(display, window, 0, ptr::null_mut());

	xlib::XMapWindow(display, window);
	xlib::XFlush(display);

	std::thread::sleep_ms(1000);

	// it is good programming practice to return system resources to the system...
	xlib::XFreeGC(display, graphics_context);
	xlib::XDestroyWindow(display, window);
	xlib::XCloseDisplay(display);

	}
}

#[test]
fn test_shm() {
	use x11::xlib;
	use std::{ ptr, mem };
	use ffi::*;
	use x11::xlib::Bool;
	use libc::{ c_uint, size_t, c_char, c_void, c_int };

	unsafe {

	// use the information from the environment variable DISPLAY to create the X connection:
	let display = xlib::XOpenDisplay(ptr::null_mut());
	if display.is_null() {
		panic!("Unable to connect X server");
	}

	println!("MIT-SHM available: {}", XShmQueryExtension(display) != 0);
	let (mut major_ver, mut minor_ver) = (0, 0);
	let mut shared_pixmaps: Bool = 0;
	XShmQueryVersion(display, &mut major_ver, &mut minor_ver, &mut shared_pixmaps);
	println!("SHM: major: {}, minor: {}, shared pixmaps: {}", major_ver, minor_ver, shared_pixmaps);

	let screen_num = xlib::XDefaultScreen(display);
	println!("Screen: {}", screen_num);

	let mut root_window = 0;
	let (mut root_x, mut root_y) = (0, 0);
	let (mut root_width, mut root_height, mut root_border_width) = (0, 0, 0);
	let mut root_pixel_depth = 0;
	if xlib::XGetGeometry(display, xlib::XRootWindow(display, screen_num),
		&mut root_window,
		&mut root_x, &mut root_y,
		&mut root_width, &mut root_height, &mut root_border_width,
		&mut root_pixel_depth) == 0
	{
		panic!("XGetGeometry failed");
	}
	println!("root window: {}; x, y: {:?}; w, h, bw: {:?}; d: {:?}",
		root_window, (root_x, root_y), (root_width, root_height, root_border_width),
		root_pixel_depth);

	let mut shm_segment_info = mem::zeroed();
	let image_ptr = XShmCreateImage(display, xlib::XDefaultVisual(display, screen_num),
		xlib::XDefaultDepth(display, screen_num) as c_uint,
		ZPixmap, ptr::null_mut(),
		&mut shm_segment_info,
		root_width, root_height);
	if image_ptr.is_null() {
		panic!("XShmCreateImage returned null pointer");
	}
	let image = &mut *image_ptr;

	shm_segment_info.shmid = shmget(IPC_PRIVATE,
		(image.bytes_per_line * image.height) as size_t,
		IPC_CREAT | 0777);
	if shm_segment_info.shmid == -1 {
		panic!("shmget failed");
	}

	shm_segment_info.shmaddr = shmat(shm_segment_info.shmid, ptr::null_mut(), 0) as *mut c_char;
	if shm_segment_info.shmaddr.is_null() {
		panic!("XShmCreateImage returned null pointer");
	}
	image.data = shm_segment_info.shmaddr;

	shm_segment_info.read_only = 0;

	if XShmAttach(display, &mut shm_segment_info) == 0 {
		panic!("XShmAttach failed");
	}

	if XShmGetImage(display, root_window, image, root_x, root_y, AllPlanes) == 0 {
		panic!("XShmGetImage failed");
	}

	println!("width {}, height {}, xoffset {}, format {}, byte_order {}, bitmap_unity {}, \
		bitmap_bit_order {}, bitmap_pad {}, depth {}, bytes_per_line {}, bits_per_pixel {}, \
		red_mask {:b}, green_mask {:b}, blue_mask {:b}",
		image.width, image.height, image.xoffset, image.format, image.byte_order,
		image.bitmap_unity, image.bitmap_bit_order, image.bitmap_pad, image.depth,
		image.bytes_per_line, image.bits_per_pixel,
		image.red_mask, image.green_mask, image.blue_mask);

	// Factor to multiply color value with for it to fit in a u8
	let (red_size, red_shift) = mask_size_and_shift(image.red_mask);
	let (green_size, green_shift) = mask_size_and_shift(image.red_mask);
	let (blue_size, blue_shift) = mask_size_and_shift(image.red_mask);
	let (red_mask, green_mask, blue_mask) = (image.red_mask, image.green_mask, image.blue_mask);
	let masked_pixel_to_rgb: Box<Fn(c_ulong) -> RGB8> =
		if red_size == 8 && green_size == 8 && blue_size == 8
	{

		Box::new(move |pixel| RGB8{
			r: ((pixel & red_mask) >> red_shift) as u8,
			g: ((pixel & green_mask) >> green_shift) as u8,
			b: ((pixel & blue_mask) >> blue_shift) as u8,
		})
	} else {
		let red_size_factor = 8.0 / red_size as f32;
		let green_size_factor = 8.0 / green_size as f32;
		let blue_size_factor = 8.0 / blue_size as f32;
		Box::new(move |pixel| RGB8{
			r: (((pixel & red_mask) >> red_shift) as f32 * red_size_factor) as u8,
			g: (((pixel & green_mask) >> green_shift) as f32 * green_size_factor) as u8,
			b: (((pixel & blue_mask) >> blue_shift) as f32 * blue_size_factor) as u8,
		})
	};

	let mut pixel_buf = Vec::with_capacity(image.width as usize * image.height as usize);
	for row in 0..image.height {
		for col in 0..image.width {
			pixel_buf.push(masked_pixel_to_rgb(xlib::XGetPixel(image_ptr, row, col)));
		}
	}

	println!("Tot color: {:?}", pixel_buf.iter()
		.fold((0, 0, 0), |(r, g, b), p| (r + p.r as u64, g + p.g as u64, b + p.b as u64));
	);

	XShmDetach(display, &mut shm_segment_info);
	xlib::XDestroyImage(image);
	shmdt(shm_segment_info.shmaddr as *mut c_void);
	shmctl(shm_segment_info.shmid, IPC_RMID, ptr::null_mut());

	xlib::XCloseDisplay(display);

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

//! Capture the screen with xlib

#![allow(dead_code, non_upper_case_globals, non_camel_case_types)]

extern crate x11;
extern crate libc;

use libc::c_ulong;

pub mod ffi;

#[derive(Debug, Eq, PartialEq)]
struct RGB8 {
	r: u8,
	g: u8,
	b: u8,
}

/// From a bitmask for a color in a pixel, calculate the color size in bits and the bitshift
fn mask_size_and_shift(mut mask: c_ulong) -> (c_ulong, u16) {
	let (mut bits, mut shift) = (0, 0);
	let mut prev = 0;
	for _ in 0..(c_ulong::max_value() as f64).log2() as u16 {
		if mask & 1 == 1 {
			bits += 1;
			prev = 1;
		} else if prev == 1 {
			break;
		} else {
			shift += 1;
		}
		mask >>= 1;
	}
	(bits, shift)
}

// #[test]
// fn test_create_window() {
// 	use x11::xlib;
// 	use std::ptr;

// 	unsafe {

// 	// use the information from the environment variable DISPLAY to create the X connection:
// 	let display = match xlib::XOpenDisplay(ptr::null_mut()) {
// 		null if null.is_null() => panic!("Unable to connect X server"),
// 		display => display,
// 	};

// 	let screen_num = xlib::XDefaultScreen(display);
// 	let (black_pixel, white_pixel) = (xlib::XBlackPixel(display, screen_num),
// 		xlib::XWhitePixel(display, screen_num));

// 	// Once the display is initialized, create the window
// 	let (x, y) = (0, 0);
// 	let (width, height, border_width) = (400, 500, 5);
// 	let (border_color, background_color) = (black_pixel, white_pixel);
// 	let window = xlib::XCreateSimpleWindow(display,
// 		xlib::XRootWindow(display, screen_num),
// 		x, y,
// 		width, height, border_width,
// 		border_color, background_color);

// 	// this routine determines which types of input are allowed in the input.
// 	xlib::XSelectInput(display, window,
// 		xlib::ExposureMask | xlib::ButtonPressMask | xlib::KeyPressMask);

// 	// create the Graphics Context
// 	let graphics_context = xlib::XCreateGC(display, window, 0, ptr::null_mut());

// 	xlib::XMapWindow(display, window);
// 	xlib::XFlush(display);

// 	std::thread::sleep_ms(1000);

// 	// it is good programming practice to return system resources to the system...
// 	xlib::XFreeGC(display, graphics_context);
// 	xlib::XDestroyWindow(display, window);
// 	xlib::XCloseDisplay(display);

// 	}
// }

#[test]
fn test_shm() {
	use x11::xlib;
	use std::{ ptr, mem };
	use ffi::*;
	use x11::xlib::Bool;
	use libc::{ c_uint, size_t, c_char, c_void, c_int };

	unsafe {

	// use the information from the environment variable DISPLAY to create the X connection:
	let display = xlib::XOpenDisplay(ptr::null_mut());
	if display.is_null() {
		panic!("Unable to connect X server");
	}

	println!("MIT-SHM available: {}", XShmQueryExtension(display) != 0);
	let (mut major_ver, mut minor_ver) = (0, 0);
	let mut shared_pixmaps: Bool = 0;
	XShmQueryVersion(display, &mut major_ver, &mut minor_ver, &mut shared_pixmaps);
	println!("SHM: major: {}, minor: {}, shared pixmaps: {}", major_ver, minor_ver, shared_pixmaps);

	let screen_num = xlib::XDefaultScreen(display);
	println!("Screen: {}", screen_num);

	let mut root_window = 0;
	let (mut root_x, mut root_y) = (0, 0);
	let (mut root_width, mut root_height, mut root_border_width) = (0, 0, 0);
	let mut root_pixel_depth = 0;
	if xlib::XGetGeometry(display, xlib::XRootWindow(display, screen_num),
		&mut root_window,
		&mut root_x, &mut root_y,
		&mut root_width, &mut root_height, &mut root_border_width,
		&mut root_pixel_depth) == 0
	{
		panic!("XGetGeometry failed");
	}
	println!("root window: {}; x, y: {:?}; w, h, bw: {:?}; d: {:?}",
		root_window, (root_x, root_y), (root_width, root_height, root_border_width),
		root_pixel_depth);

	let mut shm_segment_info = mem::zeroed();
	let image_ptr = XShmCreateImage(display, xlib::XDefaultVisual(display, screen_num),
		xlib::XDefaultDepth(display, screen_num) as c_uint,
		ZPixmap, ptr::null_mut(),
		&mut shm_segment_info,
		root_width, root_height);
	if image_ptr.is_null() {
		panic!("XShmCreateImage returned null pointer");
	}
	let image = &mut *image_ptr;

	shm_segment_info.shmid = shmget(IPC_PRIVATE,
		(image.bytes_per_line * image.height) as size_t,
		IPC_CREAT | 0777);
	if shm_segment_info.shmid == -1 {
		panic!("shmget failed");
	}

	shm_segment_info.shmaddr = shmat(shm_segment_info.shmid, ptr::null_mut(), 0) as *mut c_char;
	if shm_segment_info.shmaddr.is_null() {
		panic!("XShmCreateImage returned null pointer");
	}
	image.data = shm_segment_info.shmaddr;

	shm_segment_info.read_only = 0;

	if XShmAttach(display, &mut shm_segment_info) == 0 {
		panic!("XShmAttach failed");
	}

	if XShmGetImage(display, root_window, image, root_x, root_y, AllPlanes) == 0 {
		panic!("XShmGetImage failed");
	}

	println!("width {}, height {}, xoffset {}, format {}, byte_order {}, bitmap_unity {}, \
		bitmap_bit_order {}, bitmap_pad {}, depth {}, bytes_per_line {}, bits_per_pixel {}, \
		red_mask {:b}, green_mask {:b}, blue_mask {:b}",
		image.width, image.height, image.xoffset, image.format, image.byte_order,
		image.bitmap_unity, image.bitmap_bit_order, image.bitmap_pad, image.depth,
		image.bytes_per_line, image.bits_per_pixel,
		image.red_mask, image.green_mask, image.blue_mask);

	// Factor to multiply color value with for it to fit in a u8
	let (red_size, red_shift) = mask_size_and_shift(image.red_mask);
	let (green_size, green_shift) = mask_size_and_shift(image.red_mask);
	let (blue_size, blue_shift) = mask_size_and_shift(image.red_mask);
	if !(red_size == 8 && green_size == 8 && blue_size == 8 && image.bits_per_pixel == 32) {
		panic!("Bits per pixel is not 32 bits or Color channels are not 8 bits");
	}

	let raw_pixels = std::slice::from_raw_parts(image.data as *const u32,
		image.height as usize * image.width as usize);

	let mut pixel_buf = Vec::with_capacity(image.width as usize * image.height as usize);
	for row in 0..image.height as usize {
		for col in 0..image.width as usize {
			// println!("{:?}", (row, col));
			let pixel = raw_pixels[row * image.width as usize + col];
			pixel_buf.push(RGB8{
				r: ((pixel & image.red_mask as u32) >> red_shift) as u8,
				g: ((pixel & image.green_mask as u32) >> green_shift) as u8,
				b: ((pixel & image.blue_mask as u32) >> blue_shift) as u8,
			});
		}
	}

	println!("Tot color: {:?}", pixel_buf.iter()
		.fold((0, 0, 0), |(r, g, b), p| (r + p.r as u64, g + p.g as u64, b + p.b as u64))
	);

	XShmDetach(display, &mut shm_segment_info);
	xlib::XDestroyImage(image);
	shmdt(shm_segment_info.shmaddr as *mut c_void);
	shmctl(shm_segment_info.shmid, IPC_RMID, ptr::null_mut());

	xlib::XCloseDisplay(display);

	}
}<|endoftext|>


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

//! Support code for encoding and decoding types.

#![doc(html_logo_url = "http://www.rust-lang.org/logos/rust-logo-128x128-blk-v2.png",
       html_favicon_url = "http://www.rust-lang.org/favicon.ico",
       html_root_url = "http://doc.rust-lang.org/rustc-serialize/")]
#![cfg_attr(test, deny(warnings))]
#![cfg_attr(test, allow(unstable))]

// test harness access
#[cfg(test)]
extern crate test;

extern crate unicode;

pub use self::serialize::{Decoder, Encoder, Decodable, Encodable,
                          DecoderHelpers, EncoderHelpers};

mod serialize;
mod collection_impls;

pub mod base64;
pub mod hex;
pub mod json;

mod rustc_serialize {
    pub use serialize::*;
}
<commit_msg>Switch to using feature attributes<commit_after>// Copyright 2012-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Support code for encoding and decoding types.

#![feature(core, collections, unicode, io, std_misc, path)]
#![doc(html_logo_url = "http://www.rust-lang.org/logos/rust-logo-128x128-blk-v2.png",
       html_favicon_url = "http://www.rust-lang.org/favicon.ico",
       html_root_url = "http://doc.rust-lang.org/rustc-serialize/")]
#![cfg_attr(test, deny(warnings))]
#![cfg_attr(test, feature(test, rand, hash))]

// test harness access
#[cfg(test)]
extern crate test;

extern crate unicode;

pub use self::serialize::{Decoder, Encoder, Decodable, Encodable,
                          DecoderHelpers, EncoderHelpers};

mod serialize;
mod collection_impls;

pub mod base64;
pub mod hex;
pub mod json;

mod rustc_serialize {
    pub use serialize::*;
}
<|endoftext|>


<commit_before>extern crate num_cpus;
pub mod filter;
pub mod collection;

#[cfg(test)]
mod tests {
    use filter::{
        prime_filter,
        prime_filter_section,
        old_prime_filter,
    };
    extern crate time;
    extern crate num_cpus;
    use self::time::PreciseTime;
    #[test]
    fn test_prime_filter(){
        let n = 12345;
        let new_pf = prime_filter(n);
        let old_pf = old_prime_filter(n);
        for i in 0..n{
            assert_eq!(new_pf[i], old_pf[i], "Failed on number {}", i)
        }
    }
    use collection::{
        primes,
        primes_section,
        primes_concurrently,
        // primes_section_sequentially,
    };
    #[test]
    fn test_prime_collection(){
        let some_primes = [ 2, 3, 5, 7, 11, 13,
                            17, 19, 23, 29, 31,
                            37, 41, 43, 47, 53,
                            59, 61, 67, 71, 73,
                            79, 83, 89, 97, 101,
                            103, 107, 109, 113, 127,
                            131, 137, 139, 149, 151,
                            157, 163, 167, 173, 179,
                            181, 191, 193, 197, 199];

        let test_primes: Vec<usize> = primes(200);
        for i in 0..test_primes.len(){
            assert_eq!(test_primes[i], some_primes[i], "Mismatch")
        };
    }
    #[test]
    fn test_prime_sec_collection(){
        let some_primes = [ 101,
                            103, 107, 109, 113, 127,
                            131, 137, 139, 149, 151,
                            157, 163, 167, 173, 179,
                            181, 191, 193, 197, 199];

        let test_primes: Vec<usize> = primes_section(100, 200);
        for i in 0..test_primes.len(){
            assert_eq!(test_primes[i], some_primes[i], "Mismatch")
        };
    }
    #[test]
    fn can_take_any_section_filter(){
        let some_primes = old_prime_filter(200);
        for min in 0..200{
            for max in (min+1)..200{
                for (i, is_prime) in prime_filter_section(min, max).into_iter().enumerate(){
                    assert_eq!(some_primes[min+i], is_prime, "bad case from {} to {}", min, max);
                }
            }
        }
    }
    #[test]
    fn can_take_any_section(){
        let some_primes = old_prime_filter(200);
        for min in 0..200{
            for max in (min+1)..200{
                for prime in primes_section(min, max).into_iter(){
                    assert!(some_primes[prime], "bad case from {} to {}: prime: {}", min, max, prime);
                }
            }
        }
    }
    #[test]
    fn count_primes(){
        let n = 100_000_000;
        let start = PreciseTime::now();
        let primes = primes(n);
        let end = PreciseTime::now();
        println!("{} seconds to find all primes less than {}!", start.to(end), n);
        println!("{} primes total!", primes.len());
        let total_primes = primes.len();
        assert_eq!(total_primes, 5761455);
    }
    #[test]
    fn count_primes_concurrently(){
        let threads = num_cpus::get()*4;
        let n = 1_000_000_000;
        let start = PreciseTime::now();
        let primes = primes_concurrently(n, threads);
        let end = PreciseTime::now();
        println!("{} seconds to find all primes less than {}! (Threads: {})", start.to(end), n, threads);
        println!("{} primes total!", primes.len());
        let total_primes = primes.len();
        assert_eq!(total_primes, 50_847_534);
    }
    #[test]
    fn count_prime_section(){
        let min = 1_000_000;
        let max = 10_000_000;
        let primes = primes_section(min, max);
        let total_primes = primes.len();
        assert_eq!(total_primes, 586081);
    }

}
<commit_msg>Test takes too long...<commit_after>extern crate num_cpus;
pub mod filter;
pub mod collection;

#[cfg(test)]
mod tests {
    use filter::{
        prime_filter,
        prime_filter_section,
        old_prime_filter,
    };
    extern crate time;
    extern crate num_cpus;
    use self::time::PreciseTime;
    #[test]
    fn test_prime_filter(){
        let n = 12345;
        let new_pf = prime_filter(n);
        let old_pf = old_prime_filter(n);
        for i in 0..n{
            assert_eq!(new_pf[i], old_pf[i], "Failed on number {}", i)
        }
    }
    use collection::{
        primes,
        primes_section,
        primes_concurrently,
        // primes_section_sequentially,
    };
    #[test]
    fn test_prime_collection(){
        let some_primes = [ 2, 3, 5, 7, 11, 13,
                            17, 19, 23, 29, 31,
                            37, 41, 43, 47, 53,
                            59, 61, 67, 71, 73,
                            79, 83, 89, 97, 101,
                            103, 107, 109, 113, 127,
                            131, 137, 139, 149, 151,
                            157, 163, 167, 173, 179,
                            181, 191, 193, 197, 199];

        let test_primes: Vec<usize> = primes(200);
        for i in 0..test_primes.len(){
            assert_eq!(test_primes[i], some_primes[i], "Mismatch")
        };
    }
    #[test]
    fn test_prime_sec_collection(){
        let some_primes = [ 101,
                            103, 107, 109, 113, 127,
                            131, 137, 139, 149, 151,
                            157, 163, 167, 173, 179,
                            181, 191, 193, 197, 199];

        let test_primes: Vec<usize> = primes_section(100, 200);
        for i in 0..test_primes.len(){
            assert_eq!(test_primes[i], some_primes[i], "Mismatch")
        };
    }
    #[test]
    fn can_take_any_section_filter(){
        let some_primes = old_prime_filter(200);
        for min in 0..200{
            for max in (min+1)..200{
                for (i, is_prime) in prime_filter_section_sequentially(min, max).into_iter().enumerate(){
                    assert_eq!(some_primes[min+i], is_prime, "bad case from {} to {}", min, max);
                }
            }
        }
    }
    #[test]
    fn can_take_any_section(){
        let some_primes = old_prime_filter(200);
        for min in 0..200{
            for max in (min+1)..200{
                for prime in primes_section_sequentially(min, max).into_iter(){
                    assert!(some_primes[prime], "bad case from {} to {}: prime: {}", min, max, prime);
                }
            }
        }
    }
    #[test]
    fn count_primes(){
        let n = 100_000_000;
        let start = PreciseTime::now();
        let primes = primes(n);
        let end = PreciseTime::now();
        println!("{} seconds to find all primes less than {}!", start.to(end), n);
        println!("{} primes total!", primes.len());
        let total_primes = primes.len();
        assert_eq!(total_primes, 5761455);
    }
    #[test]
    fn count_primes_concurrently(){
        let threads = num_cpus::get()*4;
        let n = 1_000_000_000;
        let start = PreciseTime::now();
        let primes = primes_concurrently(n, threads);
        let end = PreciseTime::now();
        println!("{} seconds to find all primes less than {}! (Threads: {})", start.to(end), n, threads);
        println!("{} primes total!", primes.len());
        let total_primes = primes.len();
        assert_eq!(total_primes, 50_847_534);
    }
    #[test]
    fn count_prime_section(){
        let min = 1_000_000;
        let max = 10_000_000;
        let primes = primes_section(min, max);
        let total_primes = primes.len();
        assert_eq!(total_primes, 586081);
    }

}
<|endoftext|>


<commit_before>#![feature(test)]

#![feature(plugin)]
#![plugin(clippy)]

extern crate test;
extern crate rand;

use rand::distributions::{IndependentSample, Range};

pub struct Settings<R>
    where R: rand::Rng
{
    pub min_pos: Vec<f32>,
    pub max_pos: Vec<f32>,

    pub cr_min: f32,
    pub cr_max: f32,
    pub cr_change_probability: f32,

    pub f_min: f32,
    pub f_max: f32,
    pub f_change_probability: f32,

    pub pop_size: usize,
    pub rng: R,
}

impl<R> Settings<R>
    where R: rand::Rng
{
    pub fn min_max_rng(min_pos: Vec<f32>, max_pos: Vec<f32>, rng: R) -> Settings<R> {
        // create settings for the algorithm
        Settings {
            min_pos: min_pos,
            max_pos: max_pos,

            cr_min: 0.0,
            cr_max: 1.0,
            cr_change_probability: 0.1,

            f_min: 0.1,
            f_max: 1.0,
            f_change_probability: 0.1,

            pop_size: 50,
            rng: rng,
        }
    }
}

impl Settings<rand::XorShiftRng> {
    pub fn new(min_pos: Vec<f32>, max_pos: Vec<f32>) -> Settings<rand::XorShiftRng> {
        Settings::min_max_rng(min_pos, max_pos, rand::weak_rng())
    }
}

#[derive(Clone,Debug)]
pub struct Individual {
    pub pos: Vec<f32>,
    // the lower, the better.
    pub cost: Option<f32>,

    // control parameters
    cr: f32,
    f: f32,
}

pub struct Population<R>
    where R: rand::Rng
{
    // TODO use a single vector for curr and best and controlparameters?
    pub curr: Vec<Individual>,
    best: Vec<Individual>,

    settings: Settings<R>,
    best_idx: Option<usize>,
    dim: usize,
    between_popsize: Range<usize>,
    between_dim: Range<usize>,
    between_cr: Range<f32>,
    between_f: Range<f32>,
}

impl<R> Population<R>
    where R: rand::Rng
{
    // Creates a new population based on the given settings.
    pub fn new(s: Settings<R>) -> Population<R> {
        assert_eq!(s.min_pos.len(),
                   s.max_pos.len(),
                   "min_pos and max_pos need to have the same number of elements");
        assert!(s.min_pos.len() >= 1,
                "need at least one element to optimize");

        // create a vector of randomly initialized individuals for current.
        let dim = s.min_pos.len();

        // Empty individual, with no cost value (yet)
        let dummy_individual = Individual {
            pos: vec![0.0; dim],
            cost: None,
            cr: 0.0,
            f: 0.0,
        };

        // creates all the empty individuals
        let mut pop = Population {
            curr: vec![dummy_individual.clone(); s.pop_size],
            best: vec![dummy_individual; s.pop_size],
            best_idx: None,
            dim: dim,
            between_popsize: Range::new(0, s.pop_size),
            between_dim: Range::new(0, dim),
            between_cr: Range::new(s.cr_min, s.cr_max),
            between_f: Range::new(s.f_min, s.f_max),
            settings: s,
        };

        for ind in &mut pop.curr {
            // init control parameters
            ind.cr = pop.between_cr.ind_sample(&mut pop.settings.rng);
            ind.f = pop.between_f.ind_sample(&mut pop.settings.rng);

            // random range for each dimension
            for d in 0..dim {
                let between_min_max = Range::new(pop.settings.min_pos[d], pop.settings.max_pos[d]);
                ind.pos[d] = between_min_max.ind_sample(&mut pop.settings.rng);
            }
        }

        pop
    }

    fn update_best(&mut self) -> bool {
        let last_best_cost: Option<f32>;
        if let Some(last_best_idx) = self.best_idx {
            last_best_cost = self.best[last_best_idx].cost;
        } else {
            last_best_cost = None;
        }

        let mut new_best_cost = last_best_cost;
        for i in 0..self.curr.len() {
            let curr = &mut self.curr[i];
            let best = &mut self.best[i];

            // we use <= here, so that the individual moves even if the cost
            // stays the same.
            if best.cost.is_none() || curr.cost.unwrap() <= best.cost.unwrap() {
                // find global best. We use < here so that global only updates when fitness changes.
                if new_best_cost.is_none() || curr.cost.unwrap() < new_best_cost.unwrap() {
                    self.best_idx = Some(i);
                    new_best_cost = curr.cost;
                }

                // replace individual's best. swap is *much* faster than clone.
                std::mem::swap(curr, best);
            }
        }

        // got a new best?
        last_best_cost.is_none() || new_best_cost.unwrap() < last_best_cost.unwrap()
    }

    fn update_positions(&mut self) {
        let global_best_pos = &self.best[self.best_idx.unwrap()].pos;
        let rng = &mut self.settings.rng;
        for i in 0..self.curr.len() {
            let mut id1 = i;
            while id1 == i {
                id1 = self.between_popsize.ind_sample(rng);
            }

            let mut id2 = i;
            while id2 == i || id2 == id1 {
                id2 = self.between_popsize.ind_sample(rng);
            }

            let curr = &mut self.curr[i];
            let best = &self.best[i];

            // see "Self-Adapting Control Parameters in Differential Evolution:
            // A Comparative Study on Numerical Benchmark Problems"
            if rng.next_f32() < self.settings.cr_change_probability {
                curr.cr = self.between_cr.ind_sample(rng);
            } else {
                curr.cr = best.cr;
            }
            if rng.next_f32() < self.settings.f_change_probability {
                curr.f = self.between_f.ind_sample(rng);
            } else {
                curr.f = best.f;
            }

            let curr_pos = &mut curr.pos;
            let best_pos = &best.pos;
            let best1_pos = &self.best[id1].pos;
            let best2_pos = &self.best[id2].pos;

            let forced_mutation_dim = self.between_dim.ind_sample(rng);
            for d in 0..self.dim {
                if d == forced_mutation_dim || rng.next_f32() < curr.cr {
                    curr_pos[d] = global_best_pos[d] + curr.f * (best1_pos[d] - best2_pos[d]);
                } else {
                    curr_pos[d] = best_pos[d];
                }
            }

            // reset cost, has to be updated by the user.
            curr.cost = None;
        }
    }

    // Uses updated cost values to update positions of individuals.
    pub fn evolve(&mut self) -> Option<&Individual> {
        let found_new_best = self.update_best();
        self.update_positions();

        if found_new_best {
            Some(&self.best[self.best_idx.unwrap()])
        } else {
            None
        }
    }
}

#[cfg(test)]
mod tests {
    extern crate rand;
    use super::*;
    use test::Bencher;
    use rand::{XorShiftRng, StdRng, IsaacRng, Isaac64Rng, Rng};
    use rand::{OsRng, weak_rng};

    fn setup<R: rand::Rng>(dim: usize, rng: R) -> Population<R> {
        let settings = Settings::min_max_rng(vec![-20.0; dim], vec![20.0; dim], rng);
        Population::new(settings)
    }

    fn dummy_fitness(individuals: &mut Vec<Individual>) {
        for ind in individuals {
            ind.cost = Some(1.234);
        }
    }

    #[bench]
    fn bench_square_fitness_opt(b: &mut Bencher) {
        let rng: XorShiftRng = OsRng::new().unwrap().gen();
        let mut pop = setup(5, rng);
        b.iter(|| {
            for ind in &mut pop.curr {
                let mut f = 0.0;
                for x in &ind.pos {
                    f += x * x;
                }
                ind.cost = Some(f);
            }
            pop.evolve();
        });
    }

    #[bench]
    fn rand_thread_rng(b: &mut Bencher) {
        let mut pop = setup(5, rand::thread_rng());
        b.iter(|| {
            dummy_fitness(&mut pop.curr);
            pop.evolve();
        });
    }

    #[bench]
    fn rand_xor_shift(b: &mut Bencher) {
        let rng: XorShiftRng = OsRng::new().unwrap().gen();
        let mut pop = setup(5, rng);
        b.iter(|| {
            dummy_fitness(&mut pop.curr);
            pop.evolve();
        });
    }

    #[bench]
    fn rand_isaac(b: &mut Bencher) {
        let rng: IsaacRng = OsRng::new().unwrap().gen();
        let mut pop = setup(5, rng);
        b.iter(|| {
            dummy_fitness(&mut pop.curr);
            pop.evolve();
        });
    }

    #[bench]
    fn rand_isaac64(b: &mut Bencher) {
        let rng: Isaac64Rng = OsRng::new().unwrap().gen();
        let mut pop = setup(5, rng);
        b.iter(|| {
            dummy_fitness(&mut pop.curr);
            pop.evolve();
        });
    }

    #[bench]
    fn rand_std(b: &mut Bencher) {
        let rng: StdRng = StdRng::new().unwrap();
        let mut pop = setup(5, rng);
        b.iter(|| {
            dummy_fitness(&mut pop.curr);
            pop.evolve();
        });
    }

    #[bench]
    fn rand_weak_rng(b: &mut Bencher) {
        let rng = weak_rng();
        let mut pop = setup(5, rng);
        b.iter(|| {
            dummy_fitness(&mut pop.curr);
            pop.evolve();
        });
    }

}
<commit_msg>use gen::<f32> instead of next_32 this is as fast and more generic.<commit_after>#![feature(test)]

#![feature(plugin)]
#![plugin(clippy)]

extern crate test;
extern crate rand;

use rand::distributions::{IndependentSample, Range};

pub struct Settings<R>
    where R: rand::Rng
{
    pub min_pos: Vec<f32>,
    pub max_pos: Vec<f32>,

    pub cr_min: f32,
    pub cr_max: f32,
    pub cr_change_probability: f32,

    pub f_min: f32,
    pub f_max: f32,
    pub f_change_probability: f32,

    pub pop_size: usize,
    pub rng: R,
}

impl<R> Settings<R>
    where R: rand::Rng
{
    pub fn min_max_rng(min_pos: Vec<f32>, max_pos: Vec<f32>, rng: R) -> Settings<R> {
        // create settings for the algorithm
        Settings {
            min_pos: min_pos,
            max_pos: max_pos,

            cr_min: 0.0,
            cr_max: 1.0,
            cr_change_probability: 0.1,

            f_min: 0.1,
            f_max: 1.0,
            f_change_probability: 0.1,

            pop_size: 50,
            rng: rng,
        }
    }
}

impl Settings<rand::XorShiftRng> {
    pub fn new(min_pos: Vec<f32>, max_pos: Vec<f32>) -> Settings<rand::XorShiftRng> {
        Settings::min_max_rng(min_pos, max_pos, rand::weak_rng())
    }
}

#[derive(Clone,Debug)]
pub struct Individual {
    pub pos: Vec<f32>,
    // the lower, the better.
    pub cost: Option<f32>,

    // control parameters
    cr: f32,
    f: f32,
}

pub struct Population<R>
    where R: rand::Rng
{
    // TODO use a single vector for curr and best and controlparameters?
    pub curr: Vec<Individual>,
    best: Vec<Individual>,

    settings: Settings<R>,
    best_idx: Option<usize>,
    dim: usize,
    between_popsize: Range<usize>,
    between_dim: Range<usize>,
    between_cr: Range<f32>,
    between_f: Range<f32>,
}

impl<R> Population<R>
    where R: rand::Rng
{
    // Creates a new population based on the given settings.
    pub fn new(s: Settings<R>) -> Population<R> {
        assert_eq!(s.min_pos.len(),
                   s.max_pos.len(),
                   "min_pos and max_pos need to have the same number of elements");
        assert!(s.min_pos.len() >= 1,
                "need at least one element to optimize");

        // create a vector of randomly initialized individuals for current.
        let dim = s.min_pos.len();

        // Empty individual, with no cost value (yet)
        let dummy_individual = Individual {
            pos: vec![0.0; dim],
            cost: None,
            cr: 0.0,
            f: 0.0,
        };

        // creates all the empty individuals
        let mut pop = Population {
            curr: vec![dummy_individual.clone(); s.pop_size],
            best: vec![dummy_individual; s.pop_size],
            best_idx: None,
            dim: dim,
            between_popsize: Range::new(0, s.pop_size),
            between_dim: Range::new(0, dim),
            between_cr: Range::new(s.cr_min, s.cr_max),
            between_f: Range::new(s.f_min, s.f_max),
            settings: s,
        };

        for ind in &mut pop.curr {
            // init control parameters
            ind.cr = pop.between_cr.ind_sample(&mut pop.settings.rng);
            ind.f = pop.between_f.ind_sample(&mut pop.settings.rng);

            // random range for each dimension
            for d in 0..dim {
                let between_min_max = Range::new(pop.settings.min_pos[d], pop.settings.max_pos[d]);
                ind.pos[d] = between_min_max.ind_sample(&mut pop.settings.rng);
            }
        }

        pop
    }

    fn update_best(&mut self) -> bool {
        let last_best_cost: Option<f32>;
        if let Some(last_best_idx) = self.best_idx {
            last_best_cost = self.best[last_best_idx].cost;
        } else {
            last_best_cost = None;
        }

        let mut new_best_cost = last_best_cost;
        for i in 0..self.curr.len() {
            let curr = &mut self.curr[i];
            let best = &mut self.best[i];

            // we use <= here, so that the individual moves even if the cost
            // stays the same.
            if best.cost.is_none() || curr.cost.unwrap() <= best.cost.unwrap() {
                // find global best. We use < here so that global only updates when fitness changes.
                if new_best_cost.is_none() || curr.cost.unwrap() < new_best_cost.unwrap() {
                    self.best_idx = Some(i);
                    new_best_cost = curr.cost;
                }

                // replace individual's best. swap is *much* faster than clone.
                std::mem::swap(curr, best);
            }
        }

        // got a new best?
        last_best_cost.is_none() || new_best_cost.unwrap() < last_best_cost.unwrap()
    }

    fn update_positions(&mut self) {
        let global_best_pos = &self.best[self.best_idx.unwrap()].pos;
        let rng = &mut self.settings.rng;
        for i in 0..self.curr.len() {
            let mut id1 = i;
            while id1 == i {
                id1 = self.between_popsize.ind_sample(rng);
            }

            let mut id2 = i;
            while id2 == i || id2 == id1 {
                id2 = self.between_popsize.ind_sample(rng);
            }

            let curr = &mut self.curr[i];
            let best = &self.best[i];

            // see "Self-Adapting Control Parameters in Differential Evolution:
            // A Comparative Study on Numerical Benchmark Problems"
            if rng.gen::<f32>() < self.settings.cr_change_probability {
                curr.cr = self.between_cr.ind_sample(rng);
            } else {
                curr.cr = best.cr;
            }
            if rng.gen::<f32>() < self.settings.f_change_probability {
                curr.f = self.between_f.ind_sample(rng);
            } else {
                curr.f = best.f;
            }

            let curr_pos = &mut curr.pos;
            let best_pos = &best.pos;
            let best1_pos = &self.best[id1].pos;
            let best2_pos = &self.best[id2].pos;

            let forced_mutation_dim = self.between_dim.ind_sample(rng);
            for d in 0..self.dim {
                if d == forced_mutation_dim || rng.gen::<f32>() < curr.cr {
                    curr_pos[d] = global_best_pos[d] + curr.f * (best1_pos[d] - best2_pos[d]);
                } else {
                    curr_pos[d] = best_pos[d];
                }
            }

            // reset cost, has to be updated by the user.
            curr.cost = None;
        }
    }

    // Uses updated cost values to update positions of individuals.
    pub fn evolve(&mut self) -> Option<&Individual> {
        let found_new_best = self.update_best();
        self.update_positions();

        if found_new_best {
            Some(&self.best[self.best_idx.unwrap()])
        } else {
            None
        }
    }
}

#[cfg(test)]
mod tests {
    extern crate rand;
    use super::*;
    use test::Bencher;
    use rand::{XorShiftRng, StdRng, IsaacRng, Isaac64Rng, Rng};
    use rand::{OsRng, weak_rng};

    fn setup<R: rand::Rng>(dim: usize, rng: R) -> Population<R> {
        let settings = Settings::min_max_rng(vec![-20.0; dim], vec![20.0; dim], rng);
        Population::new(settings)
    }

    fn dummy_fitness(individuals: &mut Vec<Individual>) {
        for ind in individuals {
            ind.cost = Some(1.234);
        }
    }

    #[bench]
    fn bench_square_fitness_opt(b: &mut Bencher) {
        let rng: XorShiftRng = OsRng::new().unwrap().gen();
        let mut pop = setup(5, rng);
        b.iter(|| {
            for ind in &mut pop.curr {
                let mut f = 0.0;
                for x in &ind.pos {
                    f += x * x;
                }
                ind.cost = Some(f);
            }
            pop.evolve();
        });
    }

    #[bench]
    fn rand_thread_rng(b: &mut Bencher) {
        let mut pop = setup(5, rand::thread_rng());
        b.iter(|| {
            dummy_fitness(&mut pop.curr);
            pop.evolve();
        });
    }

    #[bench]
    fn rand_xor_shift(b: &mut Bencher) {
        let rng: XorShiftRng = OsRng::new().unwrap().gen();
        let mut pop = setup(5, rng);
        b.iter(|| {
            dummy_fitness(&mut pop.curr);
            pop.evolve();
        });
    }

    #[bench]
    fn rand_isaac(b: &mut Bencher) {
        let rng: IsaacRng = OsRng::new().unwrap().gen();
        let mut pop = setup(5, rng);
        b.iter(|| {
            dummy_fitness(&mut pop.curr);
            pop.evolve();
        });
    }

    #[bench]
    fn rand_isaac64(b: &mut Bencher) {
        let rng: Isaac64Rng = OsRng::new().unwrap().gen();
        let mut pop = setup(5, rng);
        b.iter(|| {
            dummy_fitness(&mut pop.curr);
            pop.evolve();
        });
    }

    #[bench]
    fn rand_std(b: &mut Bencher) {
        let rng: StdRng = StdRng::new().unwrap();
        let mut pop = setup(5, rng);
        b.iter(|| {
            dummy_fitness(&mut pop.curr);
            pop.evolve();
        });
    }

    #[bench]
    fn rand_weak_rng(b: &mut Bencher) {
        let rng = weak_rng();
        let mut pop = setup(5, rng);
        b.iter(|| {
            dummy_fitness(&mut pop.curr);
            pop.evolve();
        });
    }

}
<|endoftext|>