<commit_before>use chrono::UTC;

use params::*;

use rs_es::Client;
use rs_es::query::{Filter, Query};
use rs_es::units::{JsonVal, DurationUnit};
use rs_es::units::Duration as ESDuration;
use rs_es::operations::search::{Sort, SortField, Order};
use rs_es::operations::index::IndexResult;
use rs_es::operations::bulk::Action;
use rs_es::error::EsError;

extern crate searchspot;
use searchspot::terms::VectorOfTerms;

#[derive(RustcEncodable, RustcDecodable, Debug)]
pub struct Talent {
  pub id:                 u32,
  pub accepted:           bool,
  pub work_roles:         Vec<String>,
  pub work_languages:     Vec<String>,
  pub work_experience:    String,
  pub work_locations:     Vec<String>,
  pub work_authorization: String,
  pub company_ids:        Vec<u32>,
  pub batch_starts_at:    String,
  pub batch_ends_at:      String,
  pub added_to_batch_at:  String,
  pub weight:             i32,
  pub blocked_companies:  Vec<u32>
}

impl Talent {
  /// Populate the ElasticSearch index with `self`.
  pub fn index(&self, mut es: &mut Client, index: &str) -> Result<IndexResult, EsError> {
    es.index(index, "talent")
      .with_doc(&self)
      .with_id(&*format!("{}", self.id))
      .send()
  }

  #[allow(dead_code)]
  pub fn delete_all(mut es: &mut Client, index: &str) {
    let mut scan = match es.search_query()
                           .with_indexes(&[&index])
                           .with_query(&Query::build_match_all().build())
                           .scan(ESDuration::new(1, DurationUnit::Minute)) {
                              Ok(scan) => scan,
                              Err(_)   => return
                           };

    loop {
      let     page = scan.scroll(&mut es).unwrap();
      let mut hits = page.hits.hits;

      if hits.len() == 0 {
        break;
      }

      let actions = hits.drain(..)
                        .map(|hit| {
                           Action::delete(hit.id).with_index(&*index)
                                                 .with_doc_type(hit.doc_type)
                        })
                        .collect::<Vec<Action>>();
      es.bulk(&actions)
        .send()
        .unwrap();
    }

    scan.close(&mut es)
        .unwrap();
  }

  /// Query ElasticSearch on given `indexes` and `params` and return the IDs of
  /// the found talents.
  pub fn search(mut es: &mut Client, default_index: &str, params: &Map) -> Vec<u32> {
    let now   = UTC::now().to_rfc3339();
    let epoch = match params.find(&["epoch"]) {
      Some(epoch) => String::from_value(&epoch).unwrap_or(now),
      _           => now
    };

    let index: Vec<&str> = match params.find(&["index"]) {
      Some(&Value::String(ref index)) => vec![&index[..]],
      _ => vec![default_index]
    };

    let result = es.search_query()
                   .with_indexes(&*index)
                   .with_query(&Talent::search_filters(params, &*epoch))
                   .with_sort(&Talent::sorting_criteria())
                   .with_size(1000) // TODO
                   .send();

    match result {
      Ok(result) => {
        let mut results = result.hits.hits.into_iter()
                                          .map(|hit| hit.source.unwrap()["id"]
                                                               .as_u64()
                                                               .unwrap() as u32)
                                          .collect::<Vec<u32>>();
        results.dedup();
        results
      },
      Err(err) => {
        println!("{:?}", err);
        vec![]
      }
    }
  }

  /// Return a `Vec<Filter>` with visibility criteria for the talents.
  /// The `epoch` must be given as `I64` (UNIX time in seconds) and is
  /// the range in which batches are searched.
  /// If `presented_talents` is provided, talents who match the IDs
  /// contained there skip the standard visibility criteria.
  ///
  /// Basically, the talents must be accepted into the platform and must be
  /// inside a living batch to match the visibility criteria.
  fn visibility_filters(epoch: &str, presented_talents: Vec<i32>) -> Vec<Filter> {
    let visibility_rules = Filter::build_bool()
                                  .with_must(
                                    vec![
                                      Filter::build_term("accepted", true)
                                             .build(),
                                      Filter::build_range("batch_starts_at")
                                             .with_lte(JsonVal::from(epoch))
                                             .with_format("strict_date_optional_time")
                                             .build(),
                                      Filter::build_range("batch_ends_at")
                                             .with_gte(JsonVal::from(epoch))
                                             .with_format("strict_date_optional_time")
                                             .build()
                                    ])
                                  .build();

    if presented_talents.len() > 0 { // preferred over !_.is_empty()
      let presented_talents_filters = Filter::build_bool()
                                             .with_must(
                                               vec![
                                                 <Filter as VectorOfTerms<i32>>::build_terms(
                                                   "ids", &presented_talents)
                                               ].into_iter()
                                                .flat_map(|x| x)
                                                .collect::<Vec<Filter>>())
                                             .build();
      vec![
        Filter::build_bool()
               .with_should(vec![visibility_rules, presented_talents_filters])
               .build()
      ]
    }
    else {
      vec![visibility_rules]
    }
  }

  /// Given parameters inside the query string mapped inside a `Map`,
  /// and the `epoch` (defined as UNIX time in seconds) for batches,
  /// return a `Query` for ElasticSearch.
  ///
  /// `VectorOfTerms` are ORred, while `Filter`s are ANDed.
  /// I.e.: given ["Fullstack", "DevOps"] as `work_roles`, found talents
  /// will present at least one of these roles), but both `work_roles`
  /// and `work_languages`, if provided, must not be empty.
  fn search_filters(params: &Map, epoch: &str) -> Query {
    let company_id = i32_vec_from_params!(params, "company_id");

    Query::build_filtered(Filter::build_bool()
                                 .with_must(
                                   vec![
                                     <Filter as VectorOfTerms<String>>::build_terms(
                                       "work_roles", &vec_from_params!(params, "work_roles")),

                                     <Filter as VectorOfTerms<String>>::build_terms(
                                       "work_languages", &vec_from_params!(params, "work_languages")),

                                     <Filter as VectorOfTerms<String>>::build_terms(
                                       "work_experience", &string_vec_from_params!(params, "work_experience")),

                                     <Filter as VectorOfTerms<String>>::build_terms(
                                      "work_authorization", &string_vec_from_params!(params, "work_authorization")),

                                     <Filter as VectorOfTerms<String>>::build_terms(
                                       "work_locations", &vec_from_params!(params, "work_locations")),

                                     Talent::visibility_filters(epoch,
                                       i32_vec_from_params!(params, "presented_talents"))
                                   ].into_iter()
                                    .flat_map(|x| x)
                                    .collect::<Vec<Filter>>())
                                 .with_must_not(
                                   vec![
                                     <Filter as VectorOfTerms<i32>>::build_terms(
                                       "company_ids", &company_id),

                                     <Filter as VectorOfTerms<i32>>::build_terms(
                                       "blocked_companies", &company_id)
                                   ].into_iter()
                                    .flat_map(|x| x)
                                    .collect::<Vec<Filter>>())
                                 .build())
          .build()
  }

  /// Return a `Sort` that makes values be sorted for given fields, descendently.
  fn sorting_criteria() -> Sort {
    Sort::new(
      vec![
        SortField::new("batch_starts_at",   Some(Order::Desc)).build(),
        SortField::new("weight",            Some(Order::Desc)).build(),
        SortField::new("added_to_batch_at", Some(Order::Desc)).build()
      ])
  }
}

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

  use rs_es::Client;

  use params::*;

  extern crate searchspot;
  use searchspot::config::*;

  use resources::user::Talent;

  use std::time::Duration as TimeDuration;
  use std::thread::sleep;

  const CONFIG_FILE: &'static str = "examples/tests.toml";

  lazy_static! {
    static ref config: Config = Config::from_file(CONFIG_FILE.to_owned());
  }

  pub fn make_client() -> Client {
    Client::new(&*config.es.host, config.es.port)
  }

  macro_rules! epoch_from_year {
    ($year:expr) => {
      UTC.datetime_from_str(&format!("{}-01-01 12:00:00", $year),
        "%Y-%m-%d %H:%M:%S").unwrap().to_rfc3339()
    }
  }

  pub fn populate_es(mut client: &mut Client) {
    vec![
      Talent {
        id:                 1,
        accepted:           true,
        work_roles:         vec![],
        work_languages:     vec![],
        work_experience:    "1..2".to_owned(),
        work_locations:     vec!["Berlin".to_owned()],
        work_authorization: "yes".to_owned(),
        company_ids:        vec![],
        batch_starts_at:    epoch_from_year!("2006"),
        batch_ends_at:      epoch_from_year!("2020"),
        added_to_batch_at:  epoch_from_year!("2006"),
        weight:             -5,
        blocked_companies:  vec![]
      },

      Talent {
        id:                 2,
        accepted:           true,
        work_roles:         vec![],
        work_languages:     vec![],
        work_experience:    "1..2".to_owned(),
        work_locations:     vec!["Berlin".to_owned()],
        work_authorization: "yes".to_owned(),
        company_ids:        vec![],
        batch_starts_at:    epoch_from_year!("2006"),
        batch_ends_at:      epoch_from_year!("2020"),
        added_to_batch_at:  epoch_from_year!("2006"),
        weight:             6,
        blocked_companies:  vec![]
      },

      Talent {
        id:                 3,
        accepted:           false,
        work_roles:         vec![],
        work_languages:     vec![],
        work_experience:    "1..2".to_owned(),
        work_locations:     vec!["Berlin".to_owned()],
        work_authorization: "yes".to_owned(),
        company_ids:        vec![],
        batch_starts_at:    epoch_from_year!("2007"),
        batch_ends_at:      epoch_from_year!("2020"),
        added_to_batch_at:  epoch_from_year!("2011"),
        weight:             6,
        blocked_companies:  vec![]
      },

      Talent {
        id:                 4,
        accepted:           true,
        work_roles:         vec!["Fullstack".to_owned(), "DevOps".to_owned()],
        work_languages:     vec![],
        work_experience:    "1..2".to_owned(),
        work_locations:     vec!["Berlin".to_owned()],
        work_authorization:  "yes".to_owned(),
        company_ids:        vec![6],
        batch_starts_at:    epoch_from_year!("2008"),
        batch_ends_at:      epoch_from_year!("2020"),
        added_to_batch_at:  epoch_from_year!("2011"),
        weight:             0,
        blocked_companies:  vec![]
      }
    ].iter()
     .all(|user| user.index(&mut client, &config.es.index)
                     .is_ok());

    sleep(TimeDuration::from_millis(5000));
  }

  #[test]
  fn test_search() {
    let mut client = make_client();
    Talent::delete_all(&mut client, &config.es.index);
    populate_es(&mut client);

    // no parameters are given
    {
      let results = Talent::search(&mut client, &*config.es.index, &Map::new());
      assert_eq!(vec![4, 2, 1], results);
    }

    // a non existing index is given
    {
      let mut map = Map::new();
      map.assign("index", Value::String("lololol".to_owned())).unwrap();

      let results = Talent::search(&mut client, &*config.es.index, &map);
      assert!(results.is_empty());
    }

    // a date that doesn't match given indexes is given
    {
      let mut map = Map::new();
      map.assign("epoch", Value::String(epoch_from_year!("2040"))).unwrap();

      let results = Talent::search(&mut client, &*config.es.index, &map);
      assert!(results.is_empty());
    }

    // TODO: map field type properly (see benashford/rs-es#11)
    // // filtering for valid work roles
    // // {
    // //   let mut map = Map::new();
    // //   map.assign("work_roles[]", Value::String("Fullstack".to_owned())).unwrap();
    // //
    // //   let results = Talent::search(&mut client, &*config.es.index, &map);
    // //   assert_eq!(vec![4], results);
    // // }

    // filtering for given company_id
    {
      let mut map = Map::new();
      map.assign("company_id", Value::String("6".into())).unwrap();

      let results = Talent::search(&mut client, &*config.es.index, &map);
      assert_eq!(vec![2, 1], results);
    }
  }
}
<commit_msg>resources/user: implement index_many<commit_after>use chrono::UTC;
use rustc_serialize::json::{self, ToJson};

use params::*;

use rs_es::Client;
use rs_es::query::{Filter, Query};
use rs_es::units::{JsonVal, DurationUnit};
use rs_es::units::Duration as ESDuration;
use rs_es::operations::search::{Sort, SortField, Order};
use rs_es::operations::index::IndexResult;
use rs_es::operations::bulk::{Action, BulkResult};
use rs_es::error::EsError;

extern crate searchspot;
use searchspot::terms::VectorOfTerms;

#[derive(RustcEncodable, RustcDecodable, Debug)]
pub struct Talent {
  pub id:                 u32,
  pub accepted:           bool,
  pub work_roles:         Vec<String>,
  pub work_languages:     Vec<String>,
  pub work_experience:    String,
  pub work_locations:     Vec<String>,
  pub work_authorization: String,
  pub company_ids:        Vec<u32>,
  pub batch_starts_at:    String,
  pub batch_ends_at:      String,
  pub added_to_batch_at:  String,
  pub weight:             i32,
  pub blocked_companies:  Vec<u32>
}

impl Talent {
  /// Populate the ElasticSearch index with `self`.
  pub fn index(&self, mut es: &mut Client, index: &str) -> Result<IndexResult, EsError> {
    es.index(index, "talent")
      .with_doc(&self)
      .with_id(&*self.id.to_string())
      .send()
  }

  /// Populate the ElasticSearch index with `Vec<Talent>`.
  pub fn index_many(talents: Vec<Talent>, mut es: &mut Client, index: &str) -> Result<BulkResult, EsError> {
    let actions: Vec<Action> = talents.into_iter().map(|talent| {
      let json = json::encode(&talent).unwrap().to_json();
      Action::index(json).with_id(&*talent.id.to_string())
                                     .with_index(index)
    }).collect();

    es.bulk(&actions).send()
  }

  #[allow(dead_code)]
  pub fn delete_all(mut es: &mut Client, index: &str) {
    let mut scan = match es.search_query()
                           .with_indexes(&[&index])
                           .with_query(&Query::build_match_all().build())
                           .scan(ESDuration::new(1, DurationUnit::Minute)) {
                              Ok(scan) => scan,
                              Err(_)   => return
                           };

    loop {
      let     page = scan.scroll(&mut es).unwrap();
      let mut hits = page.hits.hits;

      if hits.len() == 0 {
        break;
      }

      let actions = hits.drain(..)
                        .map(|hit| {
                           Action::delete(hit.id).with_index(&*index)
                                                 .with_doc_type(hit.doc_type)
                        })
                        .collect::<Vec<Action>>();
      es.bulk(&actions)
        .send()
        .unwrap();
    }

    scan.close(&mut es)
        .unwrap();
  }

  /// Query ElasticSearch on given `indexes` and `params` and return the IDs of
  /// the found talents.
  pub fn search(mut es: &mut Client, default_index: &str, params: &Map) -> Vec<u32> {
    let now   = UTC::now().to_rfc3339();
    let epoch = match params.find(&["epoch"]) {
      Some(epoch) => String::from_value(&epoch).unwrap_or(now),
      _           => now
    };

    let index: Vec<&str> = match params.find(&["index"]) {
      Some(&Value::String(ref index)) => vec![&index[..]],
      _ => vec![default_index]
    };

    let result = es.search_query()
                   .with_indexes(&*index)
                   .with_query(&Talent::search_filters(params, &*epoch))
                   .with_sort(&Talent::sorting_criteria())
                   .with_size(1000) // TODO
                   .send();

    match result {
      Ok(result) => {
        let mut results = result.hits.hits.into_iter()
                                          .map(|hit| hit.source.unwrap()["id"]
                                                               .as_u64()
                                                               .unwrap() as u32)
                                          .collect::<Vec<u32>>();
        results.dedup();
        results
      },
      Err(err) => {
        println!("{:?}", err);
        vec![]
      }
    }
  }

  /// Return a `Vec<Filter>` with visibility criteria for the talents.
  /// The `epoch` must be given as `I64` (UNIX time in seconds) and is
  /// the range in which batches are searched.
  /// If `presented_talents` is provided, talents who match the IDs
  /// contained there skip the standard visibility criteria.
  ///
  /// Basically, the talents must be accepted into the platform and must be
  /// inside a living batch to match the visibility criteria.
  fn visibility_filters(epoch: &str, presented_talents: Vec<i32>) -> Vec<Filter> {
    let visibility_rules = Filter::build_bool()
                                  .with_must(
                                    vec![
                                      Filter::build_term("accepted", true)
                                             .build(),
                                      Filter::build_range("batch_starts_at")
                                             .with_lte(JsonVal::from(epoch))
                                             .with_format("strict_date_optional_time")
                                             .build(),
                                      Filter::build_range("batch_ends_at")
                                             .with_gte(JsonVal::from(epoch))
                                             .with_format("strict_date_optional_time")
                                             .build()
                                    ])
                                  .build();

    if presented_talents.len() > 0 { // preferred over !_.is_empty()
      let presented_talents_filters = Filter::build_bool()
                                             .with_must(
                                               vec![
                                                 <Filter as VectorOfTerms<i32>>::build_terms(
                                                   "ids", &presented_talents)
                                               ].into_iter()
                                                .flat_map(|x| x)
                                                .collect::<Vec<Filter>>())
                                             .build();
      vec![
        Filter::build_bool()
               .with_should(vec![visibility_rules, presented_talents_filters])
               .build()
      ]
    }
    else {
      vec![visibility_rules]
    }
  }

  /// Given parameters inside the query string mapped inside a `Map`,
  /// and the `epoch` (defined as UNIX time in seconds) for batches,
  /// return a `Query` for ElasticSearch.
  ///
  /// `VectorOfTerms` are ORred, while `Filter`s are ANDed.
  /// I.e.: given ["Fullstack", "DevOps"] as `work_roles`, found talents
  /// will present at least one of these roles), but both `work_roles`
  /// and `work_languages`, if provided, must not be empty.
  fn search_filters(params: &Map, epoch: &str) -> Query {
    let company_id = i32_vec_from_params!(params, "company_id");

    Query::build_filtered(Filter::build_bool()
                                 .with_must(
                                   vec![
                                     <Filter as VectorOfTerms<String>>::build_terms(
                                       "work_roles", &vec_from_params!(params, "work_roles")),

                                     <Filter as VectorOfTerms<String>>::build_terms(
                                       "work_languages", &vec_from_params!(params, "work_languages")),

                                     <Filter as VectorOfTerms<String>>::build_terms(
                                       "work_experience", &string_vec_from_params!(params, "work_experience")),

                                     <Filter as VectorOfTerms<String>>::build_terms(
                                      "work_authorization", &string_vec_from_params!(params, "work_authorization")),

                                     <Filter as VectorOfTerms<String>>::build_terms(
                                       "work_locations", &vec_from_params!(params, "work_locations")),

                                     Talent::visibility_filters(epoch,
                                       i32_vec_from_params!(params, "presented_talents"))
                                   ].into_iter()
                                    .flat_map(|x| x)
                                    .collect::<Vec<Filter>>())
                                 .with_must_not(
                                   vec![
                                     <Filter as VectorOfTerms<i32>>::build_terms(
                                       "company_ids", &company_id),

                                     <Filter as VectorOfTerms<i32>>::build_terms(
                                       "blocked_companies", &company_id)
                                   ].into_iter()
                                    .flat_map(|x| x)
                                    .collect::<Vec<Filter>>())
                                 .build())
          .build()
  }

  /// Return a `Sort` that makes values be sorted for given fields, descendently.
  fn sorting_criteria() -> Sort {
    Sort::new(
      vec![
        SortField::new("batch_starts_at",   Some(Order::Desc)).build(),
        SortField::new("weight",            Some(Order::Desc)).build(),
        SortField::new("added_to_batch_at", Some(Order::Desc)).build()
      ])
  }
}

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

  use rs_es::Client;

  use params::*;

  extern crate searchspot;
  use searchspot::config::*;

  use resources::user::Talent;

  use std::time::Duration as TimeDuration;
  use std::thread::sleep;

  const CONFIG_FILE: &'static str = "examples/tests.toml";

  lazy_static! {
    static ref config: Config = Config::from_file(CONFIG_FILE.to_owned());
  }

  pub fn make_client() -> Client {
    Client::new(&*config.es.host, config.es.port)
  }

  macro_rules! epoch_from_year {
    ($year:expr) => {
      UTC.datetime_from_str(&format!("{}-01-01 12:00:00", $year),
        "%Y-%m-%d %H:%M:%S").unwrap().to_rfc3339()
    }
  }

  pub fn populate_es(mut client: &mut Client) {
    Talent::index_many(vec![
      Talent {
        id:                 1,
        accepted:           true,
        work_roles:         vec![],
        work_languages:     vec![],
        work_experience:    "1..2".to_owned(),
        work_locations:     vec!["Berlin".to_owned()],
        work_authorization: "yes".to_owned(),
        company_ids:        vec![],
        batch_starts_at:    epoch_from_year!("2006"),
        batch_ends_at:      epoch_from_year!("2020"),
        added_to_batch_at:  epoch_from_year!("2006"),
        weight:             -5,
        blocked_companies:  vec![]
      },

      Talent {
        id:                 2,
        accepted:           true,
        work_roles:         vec![],
        work_languages:     vec![],
        work_experience:    "1..2".to_owned(),
        work_locations:     vec!["Berlin".to_owned()],
        work_authorization: "yes".to_owned(),
        company_ids:        vec![],
        batch_starts_at:    epoch_from_year!("2006"),
        batch_ends_at:      epoch_from_year!("2020"),
        added_to_batch_at:  epoch_from_year!("2006"),
        weight:             6,
        blocked_companies:  vec![]
      },

      Talent {
        id:                 3,
        accepted:           false,
        work_roles:         vec![],
        work_languages:     vec![],
        work_experience:    "1..2".to_owned(),
        work_locations:     vec!["Berlin".to_owned()],
        work_authorization: "yes".to_owned(),
        company_ids:        vec![],
        batch_starts_at:    epoch_from_year!("2007"),
        batch_ends_at:      epoch_from_year!("2020"),
        added_to_batch_at:  epoch_from_year!("2011"),
        weight:             6,
        blocked_companies:  vec![]
      },

      Talent {
        id:                 4,
        accepted:           true,
        work_roles:         vec!["Fullstack".to_owned(), "DevOps".to_owned()],
        work_languages:     vec![],
        work_experience:    "1..2".to_owned(),
        work_locations:     vec!["Berlin".to_owned()],
        work_authorization:  "yes".to_owned(),
        company_ids:        vec![6],
        batch_starts_at:    epoch_from_year!("2008"),
        batch_ends_at:      epoch_from_year!("2020"),
        added_to_batch_at:  epoch_from_year!("2011"),
        weight:             0,
        blocked_companies:  vec![]
      }
    ], &mut client, &config.es.index);

    sleep(TimeDuration::from_millis(5000));
  }

  #[test]
  fn test_search() {
    let mut client = make_client();
    Talent::delete_all(&mut client, &config.es.index);
    populate_es(&mut client);

    // no parameters are given
    {
      let results = Talent::search(&mut client, &*config.es.index, &Map::new());
      assert_eq!(vec![4, 2, 1], results);
    }

    // a non existing index is given
    {
      let mut map = Map::new();
      map.assign("index", Value::String("lololol".to_owned())).unwrap();

      let results = Talent::search(&mut client, &*config.es.index, &map);
      assert!(results.is_empty());
    }

    // a date that doesn't match given indexes is given
    {
      let mut map = Map::new();
      map.assign("epoch", Value::String(epoch_from_year!("2040"))).unwrap();

      let results = Talent::search(&mut client, &*config.es.index, &map);
      assert!(results.is_empty());
    }

    // TODO: map field type properly (see benashford/rs-es#11)
    // // filtering for valid work roles
    // // {
    // //   let mut map = Map::new();
    // //   map.assign("work_roles[]", Value::String("Fullstack".to_owned())).unwrap();
    // //
    // //   let results = Talent::search(&mut client, &*config.es.index, &map);
    // //   assert_eq!(vec![4], results);
    // // }

    // filtering for given company_id
    {
      let mut map = Map::new();
      map.assign("company_id", Value::String("6".into())).unwrap();

      let results = Talent::search(&mut client, &*config.es.index, &map);
      assert_eq!(vec![2, 1], results);
    }
  }
}
<|endoftext|>


<commit_before>// Copyright 2016 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.

#![allow(dead_code, missing_docs, bad_style)]

use io::{self, ErrorKind};

pub mod args;
#[cfg(feature = "backtrace")]
pub mod backtrace;
pub mod condvar;
pub mod env;
pub mod ext;
pub mod fast_thread_local;
pub mod fd;
pub mod fs;
pub mod memchr;
pub mod mutex;
pub mod net;
pub mod os;
pub mod os_str;
pub mod path;
pub mod pipe;
pub mod process;
pub mod rand;
pub mod rwlock;
pub mod stack_overflow;
pub mod stdio;
pub mod syscall;
pub mod thread;
pub mod thread_local;
pub mod time;

#[cfg(not(test))]
pub fn init() {
    use alloc::oom;

    oom::set_oom_handler(oom_handler);

    // A nicer handler for out-of-memory situations than the default one. This
    // one prints a message to stderr before aborting. It is critical that this
    // code does not allocate any memory since we are in an OOM situation. Any
    // errors are ignored while printing since there's nothing we can do about
    // them and we are about to exit anyways.
    fn oom_handler() -> ! {
        use intrinsics;
        let msg = "fatal runtime error: out of memory\n";
        unsafe {
            let _ = syscall::write(2, msg.as_bytes());
            intrinsics::abort();
        }
    }
}

pub fn decode_error_kind(errno: i32) -> ErrorKind {
    match errno {
        syscall::ECONNREFUSED => ErrorKind::ConnectionRefused,
        syscall::ECONNRESET => ErrorKind::ConnectionReset,
        syscall::EPERM | syscall::EACCES => ErrorKind::PermissionDenied,
        syscall::EPIPE => ErrorKind::BrokenPipe,
        syscall::ENOTCONN => ErrorKind::NotConnected,
        syscall::ECONNABORTED => ErrorKind::ConnectionAborted,
        syscall::EADDRNOTAVAIL => ErrorKind::AddrNotAvailable,
        syscall::EADDRINUSE => ErrorKind::AddrInUse,
        syscall::ENOENT => ErrorKind::NotFound,
        syscall::EINTR => ErrorKind::Interrupted,
        syscall::EINVAL => ErrorKind::InvalidInput,
        syscall::ETIMEDOUT => ErrorKind::TimedOut,
        syscall::EEXIST => ErrorKind::AlreadyExists,

        // These two constants can have the same value on some systems,
        // but different values on others, so we can't use a match
        // clause
        x if x == syscall::EAGAIN || x == syscall::EWOULDBLOCK =>
            ErrorKind::WouldBlock,

        _ => ErrorKind::Other,
    }
}

pub fn cvt(result: Result<usize, syscall::Error>) -> io::Result<usize> {
    result.map_err(|err| io::Error::from_raw_os_error(err.errno))
}

/// On Redox, use an illegal instruction to abort
pub unsafe fn abort_internal() -> ! {
    ::core::intrinsics::abort();
}
<commit_msg>Update mod.rs<commit_after>// Copyright 2016 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.

#![allow(dead_code, missing_docs, bad_style)]

use io::{self, ErrorKind};

pub mod args;
#[cfg(feature = "backtrace")]
pub mod backtrace;
pub mod condvar;
pub mod env;
pub mod ext;
pub mod fast_thread_local;
pub mod fd;
pub mod fs;
pub mod memchr;
pub mod mutex;
pub mod net;
pub mod os;
pub mod os_str;
pub mod path;
pub mod pipe;
pub mod process;
pub mod rand;
pub mod rwlock;
pub mod stack_overflow;
pub mod stdio;
pub mod syscall;
pub mod thread;
pub mod thread_local;
pub mod time;

#[cfg(not(test))]
pub fn init() {

}

pub fn decode_error_kind(errno: i32) -> ErrorKind {
    match errno {
        syscall::ECONNREFUSED => ErrorKind::ConnectionRefused,
        syscall::ECONNRESET => ErrorKind::ConnectionReset,
        syscall::EPERM | syscall::EACCES => ErrorKind::PermissionDenied,
        syscall::EPIPE => ErrorKind::BrokenPipe,
        syscall::ENOTCONN => ErrorKind::NotConnected,
        syscall::ECONNABORTED => ErrorKind::ConnectionAborted,
        syscall::EADDRNOTAVAIL => ErrorKind::AddrNotAvailable,
        syscall::EADDRINUSE => ErrorKind::AddrInUse,
        syscall::ENOENT => ErrorKind::NotFound,
        syscall::EINTR => ErrorKind::Interrupted,
        syscall::EINVAL => ErrorKind::InvalidInput,
        syscall::ETIMEDOUT => ErrorKind::TimedOut,
        syscall::EEXIST => ErrorKind::AlreadyExists,

        // These two constants can have the same value on some systems,
        // but different values on others, so we can't use a match
        // clause
        x if x == syscall::EAGAIN || x == syscall::EWOULDBLOCK =>
            ErrorKind::WouldBlock,

        _ => ErrorKind::Other,
    }
}

pub fn cvt(result: Result<usize, syscall::Error>) -> io::Result<usize> {
    result.map_err(|err| io::Error::from_raw_os_error(err.errno))
}

/// On Redox, use an illegal instruction to abort
pub unsafe fn abort_internal() -> ! {
    ::core::intrinsics::abort();
}
<|endoftext|>


<commit_before>extern crate serde_json;
extern crate tempdir;

use std::path::PathBuf;
use std::sync::{
    mpsc::{channel, Receiver, Sender},
    Arc,
};
use std::thread;

use task_maker_cache::Cache;
use task_maker_dag::CacheMode;
use task_maker_exec::{executors::LocalExecutor, ExecutorClient};
use task_maker_format::ui::{UIMessage, UIExecutionStatus};
use task_maker_format::{ioi, EvaluationConfig, EvaluationData, TaskFormat, UISender};
use task_maker_store::*;

use turingarena::evaluation::{mem::*, record};
use turingarena::award::{Score, AwardName};
use turingarena::submission::mem::Submission;

use turingarena::evaluation::Event;
use turingarena::content::TextVariant;
use turingarena::rusage::{TimeUsage, MemoryUsage};

pub fn run_evaluation(task_path: PathBuf, submission: Submission) -> Receiver<Event> {
    let (event_tx, event_rx) = channel();

    let store_path = task_path.join(".task-maker-files");
    let file_store =
        Arc::new(FileStore::new(store_path.join("store")).expect("Cannot create the file store"));

    let (tx, rx_remote) = channel();
    let (tx_remote, rx) = channel();

    let server_file_store = file_store.clone();
    thread::Builder::new()
        .name("Executor thread".into())
        .spawn(move || {
            let num_cores = num_cpus::get();
            let sandbox_path = store_path.join("sandboxes");
            let cache = Cache::new(store_path.join("cache")).expect("Cannot create the cache");
            let executor = LocalExecutor::new(server_file_store, num_cores, sandbox_path);

            executor.evaluate(tx_remote, rx_remote, cache).unwrap();
        })
        .expect("Failed to spawn the executor thread");

    thread::Builder::new()
        .name("Client thread".into())
        .spawn(move || {
            let dir = tempdir::TempDir::new("evaluation").unwrap();
            let solution_file = &submission.field_values[0].file;
            let solution_path = dir.path().join(&solution_file.name.0);
            std::fs::write(&solution_path, &solution_file.content)
                .expect("Unable to write solution file");
            println!("=={:?} start", solution_path);

            let (mut eval, receiver) = EvaluationData::new();

            let config = eval.dag.config_mut();

            thread::Builder::new()
                .name("Forwarder thread".into())
                .spawn(move || {
                    for m in receiver.iter() {
                        ui_message_to_events(m, &event_tx).expect("Failed to send event");
                    }
                })
                .expect("Failed to spawn the forwarder thread");

            config
                .keep_sandboxes(false)
                .dry_run(true)
                .cache_mode(CacheMode::Everything)
                .copy_exe(false)
                .extra_time(0.0);

            let eval_config = EvaluationConfig {
                solution_filter: vec![],
                booklet_solutions: false,
                solution_paths: vec![solution_path.to_owned()],
            };
            let task = ioi::Task::new(task_path, &eval_config).expect("Invalid task");
            task.execute(&mut eval, &eval_config)
                .expect("Failed to build the DAG");

            trace!("The DAG is: {:#?}", eval.dag);

            let ui_sender = eval.sender.clone();
            ExecutorClient::evaluate(eval.dag, tx, &rx, file_store, move |status| {
                ui_sender.send(UIMessage::ServerStatus { status })
            })
                .expect("Client failed");
        })
        .expect("Failed to spawn the executor thread");
    event_rx
}

fn ui_message_to_events(ui_message: UIMessage, tx: &Sender<Event>) -> Result<(), failure::Error> {
    match ui_message {
        UIMessage::IOITestcaseScore {
            subtask: _,
            testcase,
            solution: _,
            score,
            message,
        } => {
            tx.send(Event::Value(ValueEvent {
                key: record::Key(format!("testcase.{}.score", testcase)),
                value: record::Value::Score(record::ScoreValue {
                    score: Score(score as f64),
                }),
            }))?;
            tx.send(Event::Value(ValueEvent {
                key: record::Key(format!("testcase.{}.score", testcase)),
                value: record::Value::Message(record::TextValue {
                    text: vec![TextVariant {
                        value: message,
                        attributes: vec![],
                    }]
                }),
            }))?;
        }
        UIMessage::IOISubtaskScore {
            subtask,
            solution: _,
            score,
            normalized_score,
        } => {
            tx.send(Event::Value(ValueEvent {
                key: record::Key(format!("subtask.{}.score", subtask)),
                value: record::Value::Score(record::ScoreValue {
                    score: Score(score),
                }),
            }))?;
            tx.send(Event::Score(ScoreEvent {
                award_name: AwardName(format!("subtask.{}", subtask)),
                score: Score(score),
            }))?;
            tx.send(Event::Badge(BadgeEvent {
                award_name: AwardName(format!("subtask.{}", subtask)),
                badge: normalized_score > 0.0,
            }))?;
        }
        UIMessage::IOIEvaluation {
            testcase,
            subtask: _,
            solution: _,
            status,
        } => {
            if let UIExecutionStatus::Done {
                result
            } = status {
                tx.send(Event::Value(ValueEvent {
                    key: record::Key(format!("testcase.{}.time_usage", testcase)),
                    value: record::Value::TimeUsage(record::TimeUsageValue {
                        time_usage: TimeUsage(result.resources.cpu_time),
                    }),
                }))?;
                tx.send(Event::Value(ValueEvent {
                    key: record::Key(format!("testcase.{}.memory_usage", testcase)),
                    value: record::Value::MemoryUsage(record::MemoryUsageValue {
                        memory_usage: MemoryUsage(result.resources.memory as i32),
                    }),
                }))?;
            }
        }
        _ => (),
    };
    Ok(())
}
<commit_msg>task-maker: fix wrong key in evaluation event<commit_after>extern crate serde_json;
extern crate tempdir;

use std::path::PathBuf;
use std::sync::{
    mpsc::{channel, Receiver, Sender},
    Arc,
};
use std::thread;

use task_maker_cache::Cache;
use task_maker_dag::CacheMode;
use task_maker_exec::{executors::LocalExecutor, ExecutorClient};
use task_maker_format::ui::{UIMessage, UIExecutionStatus};
use task_maker_format::{ioi, EvaluationConfig, EvaluationData, TaskFormat, UISender};
use task_maker_store::*;

use turingarena::evaluation::{mem::*, record};
use turingarena::award::{Score, AwardName};
use turingarena::submission::mem::Submission;

use turingarena::evaluation::Event;
use turingarena::content::TextVariant;
use turingarena::rusage::{TimeUsage, MemoryUsage};

pub fn run_evaluation(task_path: PathBuf, submission: Submission) -> Receiver<Event> {
    let (event_tx, event_rx) = channel();

    let store_path = task_path.join(".task-maker-files");
    let file_store =
        Arc::new(FileStore::new(store_path.join("store")).expect("Cannot create the file store"));

    let (tx, rx_remote) = channel();
    let (tx_remote, rx) = channel();

    let server_file_store = file_store.clone();
    thread::Builder::new()
        .name("Executor thread".into())
        .spawn(move || {
            let num_cores = num_cpus::get();
            let sandbox_path = store_path.join("sandboxes");
            let cache = Cache::new(store_path.join("cache")).expect("Cannot create the cache");
            let executor = LocalExecutor::new(server_file_store, num_cores, sandbox_path);

            executor.evaluate(tx_remote, rx_remote, cache).unwrap();
        })
        .expect("Failed to spawn the executor thread");

    thread::Builder::new()
        .name("Client thread".into())
        .spawn(move || {
            let dir = tempdir::TempDir::new("evaluation").unwrap();
            let solution_file = &submission.field_values[0].file;
            let solution_path = dir.path().join(&solution_file.name.0);
            std::fs::write(&solution_path, &solution_file.content)
                .expect("Unable to write solution file");
            println!("=={:?} start", solution_path);

            let (mut eval, receiver) = EvaluationData::new();

            let config = eval.dag.config_mut();

            thread::Builder::new()
                .name("Forwarder thread".into())
                .spawn(move || {
                    for m in receiver.iter() {
                        ui_message_to_events(m, &event_tx).expect("Failed to send event");
                    }
                })
                .expect("Failed to spawn the forwarder thread");

            config
                .keep_sandboxes(false)
                .dry_run(true)
                .cache_mode(CacheMode::Everything)
                .copy_exe(false)
                .extra_time(0.0);

            let eval_config = EvaluationConfig {
                solution_filter: vec![],
                booklet_solutions: false,
                solution_paths: vec![solution_path.to_owned()],
            };
            let task = ioi::Task::new(task_path, &eval_config).expect("Invalid task");
            task.execute(&mut eval, &eval_config)
                .expect("Failed to build the DAG");

            trace!("The DAG is: {:#?}", eval.dag);

            let ui_sender = eval.sender.clone();
            ExecutorClient::evaluate(eval.dag, tx, &rx, file_store, move |status| {
                ui_sender.send(UIMessage::ServerStatus { status })
            })
                .expect("Client failed");
        })
        .expect("Failed to spawn the executor thread");
    event_rx
}

fn ui_message_to_events(ui_message: UIMessage, tx: &Sender<Event>) -> Result<(), failure::Error> {
    match ui_message {
        UIMessage::IOITestcaseScore {
            subtask: _,
            testcase,
            solution: _,
            score,
            message,
        } => {
            tx.send(Event::Value(ValueEvent {
                key: record::Key(format!("testcase.{}.score", testcase)),
                value: record::Value::Score(record::ScoreValue {
                    score: Score(score as f64),
                }),
            }))?;
            tx.send(Event::Value(ValueEvent {
                key: record::Key(format!("testcase.{}.message", testcase)),
                value: record::Value::Message(record::TextValue {
                    text: vec![TextVariant {
                        value: message,
                        attributes: vec![],
                    }]
                }),
            }))?;
        }
        UIMessage::IOISubtaskScore {
            subtask,
            solution: _,
            score,
            normalized_score,
        } => {
            tx.send(Event::Value(ValueEvent {
                key: record::Key(format!("subtask.{}.score", subtask)),
                value: record::Value::Score(record::ScoreValue {
                    score: Score(score),
                }),
            }))?;
            tx.send(Event::Score(ScoreEvent {
                award_name: AwardName(format!("subtask.{}", subtask)),
                score: Score(score),
            }))?;
            tx.send(Event::Badge(BadgeEvent {
                award_name: AwardName(format!("subtask.{}", subtask)),
                badge: normalized_score > 0.0,
            }))?;
        }
        UIMessage::IOIEvaluation {
            testcase,
            subtask: _,
            solution: _,
            status,
        } => {
            if let UIExecutionStatus::Done {
                result
            } = status {
                tx.send(Event::Value(ValueEvent {
                    key: record::Key(format!("testcase.{}.time_usage", testcase)),
                    value: record::Value::TimeUsage(record::TimeUsageValue {
                        time_usage: TimeUsage(result.resources.cpu_time),
                    }),
                }))?;
                tx.send(Event::Value(ValueEvent {
                    key: record::Key(format!("testcase.{}.memory_usage", testcase)),
                    value: record::Value::MemoryUsage(record::MemoryUsageValue {
                        memory_usage: MemoryUsage(result.resources.memory as i32),
                    }),
                }))?;
            }
        }
        _ => (),
    };
    Ok(())
}
<|endoftext|>


<commit_before>use geometry::Point;
use geometry::Rect;
use geometry::Size;

/**
 Elements that may be contained by a quadtree node.
 Either child nodes, a single rect, or nothing.
 */
enum Elements {
    /// Children are top left, top right, bottom right, and bottom left, respectively.
    Children(~QuadTree, ~QuadTree, ~QuadTree, ~QuadTree),
    /// A single rectangle.
    Member(Rect),
    /// Nothing.
    NoElements,
}

/**
 A quadtree node that can contain either one rectangle,
 or exactly four child nodes.
 */
pub struct QuadTree {
    rect: Rect,
    elements: Elements,
}

impl QuadTree {
    /**
     Create an empty quadtree with a zero-sized root node.
     */
    pub fn new() -> QuadTree {
        let origin = Point { x: 0 as f32, y: 0 as f32 };
        let size = Size { width: 0 as f32, height: 0 as f32 };
        let nodeRect = Rect { origin: origin, size: size };
        let node = QuadTree { rect: nodeRect, elements: NoElements };

        node
    }

    /**
     Create a quadtree with a root node with the same origin and a square
     size with side length matching the longer dimension of `rect`.
     */
    pub fn newAutosized(rect: Rect) -> QuadTree {
        let largerDimen = if rect.size.width > rect.size.height {
            rect.size.width
        } else {
            rect.size.height
        };

        let square = Rect { origin: rect.origin, size: Size { width: largerDimen, height: largerDimen } };

        QuadTree { rect: square, elements: Member(rect), }
    }

    /**
     Create a quadtree with a root node with the given origin, size, and member rectangle.
     If `insertRect` is not entirely contained within the rectangle defined by `origin` and `size`,
     higher level nodes will be created to form a quadtree with a child of the rect specified
     and other children such that the rect is in the tree.
     */
    pub fn newWithSize(origin: Point, size: Size, insertRect: Option<Rect>) -> QuadTree {
        let nodeRect = Rect { origin: origin, size: size };
        let node = QuadTree { rect: nodeRect, elements: NoElements };

        let tree =
        match insertRect {
            Some(rect) => {
                let (success, tree) = node.insertRect(rect);
                assert!(success, "Failed inserting a node into an empty quadtree.");

                tree
            }
            None => node,
        };

        tree
    }

    /**
     Insert a rectangle into the quadtree. If `toInsert` overlaps another rectangle
     already in the tree, the return value will be (false, self).
     If the root node is zero-sized, the resulting tree will have a square root node
     large enough to hold `toInsert`.
     */
    pub fn insertRect(self, toInsert: Rect) -> (bool, QuadTree) {
        if self.rect.size.width == 0.0 {
            return (true, QuadTree::newAutosized(toInsert))
        }

        fail!("insertRect not implemented for non-zero sized trees!");

        (false, self)
    }

    /**
     Find all of the rects in `self`, or its children, that are in nodes
     intersected by the given rect.
     */
    pub fn rectsInChildNodesIntersectedByRect(&self, testRect: &Rect) -> ~[Rect] {
        // If the test rect doesn't intersect us, then it can't intersect
        // any rects that we have.
        match self.rect.intersect(testRect) {
            Some(intersection) => {
                let mut rects = ~[];
                let mut nodesToCheck = ~[self];

                while nodesToCheck.len() > 0 {
                    let mut newNodesToCheck = ~[];

                    for node in nodesToCheck.iter() {
                        match node.elements {
                            Children(~ref tl, ~ref tr, ~ref br, ~ref bl) => {
                                let mut check = ~[];
                                let intersection = &intersection;
                                if tl.rect.intersects(intersection) {
                                    check.push(tl);
                                }
                                if tr.rect.intersects(intersection) {
                                    check.push(tr);
                                }
                                if br.rect.intersects(intersection) {
                                    check.push(br);
                                }
                                if bl.rect.intersects(intersection) {
                                    check.push(bl);
                                }
                                newNodesToCheck.push_all(check);
                            }
                            Member(memberRect) => rects.push(memberRect),
                            NoElements => ()
                        };
                    }

                    nodesToCheck = newNodesToCheck;
                }

                rects
            },
            None => ~[],
        }
    }
}<commit_msg>Call newWithSize from newAutosized.<commit_after>use geometry::Point;
use geometry::Rect;
use geometry::Size;

/**
 Elements that may be contained by a quadtree node.
 Either child nodes, a single rect, or nothing.
 */
enum Elements {
    /// Children are top left, top right, bottom right, and bottom left, respectively.
    Children(~QuadTree, ~QuadTree, ~QuadTree, ~QuadTree),
    /// A single rectangle.
    Member(Rect),
    /// Nothing.
    NoElements,
}

/**
 A quadtree node that can contain either one rectangle,
 or exactly four child nodes.
 */
pub struct QuadTree {
    rect: Rect,
    elements: Elements,
}

impl QuadTree {
    /**
     Create an empty quadtree with a zero-sized root node.
     */
    pub fn new() -> QuadTree {
        let origin = Point { x: 0 as f32, y: 0 as f32 };
        let size = Size { width: 0 as f32, height: 0 as f32 };
        let nodeRect = Rect { origin: origin, size: size };
        let node = QuadTree { rect: nodeRect, elements: NoElements };

        node
    }

    /**
     Create a quadtree with a root node with the same origin and a square
     size with side length matching the longer dimension of `rect`.
     */
    pub fn newAutosized(rect: Rect) -> QuadTree {
        let largerDimen = if rect.size.width > rect.size.height {
            rect.size.width
        } else {
            rect.size.height
        };

        let size = Size { width: largerDimen, height: largerDimen };

        QuadTree::newWithSize(rect.origin, size, Some(rect))
    }

    /**
     Create a quadtree with a root node with the given origin, size, and member rectangle.
     If `insertRect` is not entirely contained within the rectangle defined by `origin` and `size`,
     higher level nodes will be created to form a quadtree with a child of the rect specified
     and other children such that the rect is in the tree.
     */
    pub fn newWithSize(origin: Point, size: Size, insertRect: Option<Rect>) -> QuadTree {
        let nodeRect = Rect { origin: origin, size: size };
        let node = QuadTree { rect: nodeRect, elements: NoElements };

        let tree =
        match insertRect {
            Some(rect) => {
                let (success, tree) = node.insertRect(rect);
                assert!(success, "Failed inserting a node into an empty quadtree.");

                tree
            }
            None => node,
        };

        tree
    }

    /**
     Insert a rectangle into the quadtree. If `toInsert` overlaps another rectangle
     already in the tree, the return value will be (false, self).
     If the root node is zero-sized, the resulting tree will have a square root node
     large enough to hold `toInsert`.
     */
    pub fn insertRect(self, toInsert: Rect) -> (bool, QuadTree) {
        if self.rect.size.width == 0.0 {
            return (true, QuadTree::newAutosized(toInsert))
        }

        fail!("insertRect not implemented for non-zero sized trees!");

        (false, self)
    }

    /**
     Find all of the rects in `self`, or its children, that are in nodes
     intersected by the given rect.
     */
    pub fn rectsInChildNodesIntersectedByRect(&self, testRect: &Rect) -> ~[Rect] {
        // If the test rect doesn't intersect us, then it can't intersect
        // any rects that we have.
        match self.rect.intersect(testRect) {
            Some(intersection) => {
                let mut rects = ~[];
                let mut nodesToCheck = ~[self];

                while nodesToCheck.len() > 0 {
                    let mut newNodesToCheck = ~[];

                    for node in nodesToCheck.iter() {
                        match node.elements {
                            Children(~ref tl, ~ref tr, ~ref br, ~ref bl) => {
                                let mut check = ~[];
                                let intersection = &intersection;
                                if tl.rect.intersects(intersection) {
                                    check.push(tl);
                                }
                                if tr.rect.intersects(intersection) {
                                    check.push(tr);
                                }
                                if br.rect.intersects(intersection) {
                                    check.push(br);
                                }
                                if bl.rect.intersects(intersection) {
                                    check.push(bl);
                                }
                                newNodesToCheck.push_all(check);
                            }
                            Member(memberRect) => rects.push(memberRect),
                            NoElements => ()
                        };
                    }

                    nodesToCheck = newNodesToCheck;
                }

                rects
            },
            None => ~[],
        }
    }
}<|endoftext|>


<commit_before>//! Functions to read and write control registers.

pub use super::model_specific::{Efer, EferFlags};

use structures::paging::PhysFrame;
use PhysAddr;

/// Various control flags modifying the basic operation of the CPU.
pub struct Cr0;

impl Cr0 {
    /// Read the current set of CR0 flags.
    pub fn read() -> Cr0Flags {
        Cr0Flags::from_bits_truncate(Self::read_raw())
    }

    /// Read the current raw CR0 value.
    pub fn read_raw() -> u64 {
        let value: u64;
        unsafe {
            asm!("mov %cr0, $0" : "=r" (value));
        }
        value
    }

    /// Write CR0 flags.
    ///
    /// Preserves the value of reserved fields. Unsafe because it's possible to violate memory
    /// safety by e.g. disabling paging.
    pub unsafe fn write(flags: Cr0Flags) {
        let mut value = Self::read_raw();
        value |= flags.bits();

        asm!("mov $0, %cr0" :: "r" (value) : "memory")
    }

    /// Updates CR0 flags.
    ///
    /// Preserves the value of reserved fields. Unsafe because it's possible to violate memory
    /// safety by e.g. disabling paging.
    pub unsafe fn update<F>(f: F)
    where
        F: FnOnce(&mut Cr0Flags),
    {
        let mut flags = Self::read();
        f(&mut flags);
        Self::write(flags);
    }
}

bitflags! {
    pub struct Cr0Flags: u64 {
        const PROTECTED_MODE_ENABLE = 1 << 0;
        const MONITOR_COPROCESSOR = 1 << 1;
        const EMULATION = 1 << 2;
        const TASK_SWITCHED = 1 << 3;
        const EXTENSION_TYPE = 1 << 4;
        const NUMERIC_ERROR = 1 << 5;
        const WRITE_PROTECT = 1 << 16;
        const ALIGNMENT_MASK = 1 << 18;
        const NON_WRITE_THROUGH = 1 << 29;
        const CACHE_DISABLE = 1 << 30;
        const PAGING = 1 << 31;
    }
}

/// Contains the physical address of the level 4 page table.
pub struct Cr3;

impl Cr3 {
    /// Read the current P4 table address from the CR3 register.
    pub fn read() -> (PhysFrame, Cr3Flags) {
        let value: u64;
        unsafe {
            asm!("mov %cr3, $0" : "=r" (value));
        }
        let flags = Cr3Flags::from_bits_truncate(value);
        let addr = PhysAddr::new(value & 0x_000f_ffff_ffff_f000);
        let frame = PhysFrame::containing_address(addr);
        (frame, flags)
    }

    /// Write a new P4 table address into the CR3 register.
    ///
    /// ## Safety
    /// Changing the level 4 page table is unsafe, because it's possible to violate memory safety by
    /// changing the page mapping.
    pub unsafe fn write(frame: PhysFrame, flags: Cr3Flags) {
        let addr = frame.start_address();
        let value = addr.as_u64() | flags.bits();
        asm!("mov $0, %cr3" :: "r" (value) : "memory")
    }
}

bitflags! {
    /// Controls cache settings for the level 4 page table.
    pub struct Cr3Flags: u64 {
        /// Use a writethrough cache policy for the P4 table (else a writeback policy is used).
        const PAGE_LEVEL_WRITETHROUGH = 1 << 3;
        /// Disable caching for the P4 table.
        const PAGE_LEVEL_CACHE_DISABLE = 1 << 4;
    }
}
<commit_msg>Fix: Only OR reserved bits in Cr0::write<commit_after>//! Functions to read and write control registers.

pub use super::model_specific::{Efer, EferFlags};

use structures::paging::PhysFrame;
use PhysAddr;

/// Various control flags modifying the basic operation of the CPU.
pub struct Cr0;

impl Cr0 {
    /// Read the current set of CR0 flags.
    pub fn read() -> Cr0Flags {
        Cr0Flags::from_bits_truncate(Self::read_raw())
    }

    /// Read the current raw CR0 value.
    pub fn read_raw() -> u64 {
        let value: u64;
        unsafe {
            asm!("mov %cr0, $0" : "=r" (value));
        }
        value
    }

    /// Write CR0 flags.
    ///
    /// Preserves the value of reserved fields. Unsafe because it's possible to violate memory
    /// safety by e.g. disabling paging.
    pub unsafe fn write(flags: Cr0Flags) {
        let old_value = Self::read_raw();
        let reserved = old_value & !(Cr0Flags::all().bits());
        let new_value = reserved | flags.bits();

        asm!("mov $0, %cr0" :: "r" (new_value) : "memory")
    }

    /// Updates CR0 flags.
    ///
    /// Preserves the value of reserved fields. Unsafe because it's possible to violate memory
    /// safety by e.g. disabling paging.
    pub unsafe fn update<F>(f: F)
    where
        F: FnOnce(&mut Cr0Flags),
    {
        let mut flags = Self::read();
        f(&mut flags);
        Self::write(flags);
    }
}

bitflags! {
    pub struct Cr0Flags: u64 {
        const PROTECTED_MODE_ENABLE = 1 << 0;
        const MONITOR_COPROCESSOR = 1 << 1;
        const EMULATION = 1 << 2;
        const TASK_SWITCHED = 1 << 3;
        const EXTENSION_TYPE = 1 << 4;
        const NUMERIC_ERROR = 1 << 5;
        const WRITE_PROTECT = 1 << 16;
        const ALIGNMENT_MASK = 1 << 18;
        const NON_WRITE_THROUGH = 1 << 29;
        const CACHE_DISABLE = 1 << 30;
        const PAGING = 1 << 31;
    }
}

/// Contains the physical address of the level 4 page table.
pub struct Cr3;

impl Cr3 {
    /// Read the current P4 table address from the CR3 register.
    pub fn read() -> (PhysFrame, Cr3Flags) {
        let value: u64;
        unsafe {
            asm!("mov %cr3, $0" : "=r" (value));
        }
        let flags = Cr3Flags::from_bits_truncate(value);
        let addr = PhysAddr::new(value & 0x_000f_ffff_ffff_f000);
        let frame = PhysFrame::containing_address(addr);
        (frame, flags)
    }

    /// Write a new P4 table address into the CR3 register.
    ///
    /// ## Safety
    /// Changing the level 4 page table is unsafe, because it's possible to violate memory safety by
    /// changing the page mapping.
    pub unsafe fn write(frame: PhysFrame, flags: Cr3Flags) {
        let addr = frame.start_address();
        let value = addr.as_u64() | flags.bits();
        asm!("mov $0, %cr3" :: "r" (value) : "memory")
    }
}

bitflags! {
    /// Controls cache settings for the level 4 page table.
    pub struct Cr3Flags: u64 {
        /// Use a writethrough cache policy for the P4 table (else a writeback policy is used).
        const PAGE_LEVEL_WRITETHROUGH = 1 << 3;
        /// Disable caching for the P4 table.
        const PAGE_LEVEL_CACHE_DISABLE = 1 << 4;
    }
}
<|endoftext|>


<commit_before>//! Functions to read and write control registers.

pub use super::model_specific::{Efer, EferFlags};

use bitflags::bitflags;

/// Various control flags modifying the basic operation of the CPU.
#[derive(Debug)]
pub struct Cr0;

bitflags! {
    /// Configuration flags of the Cr0 register.
    pub struct Cr0Flags: u64 {
        /// Enables protected mode.
        const PROTECTED_MODE_ENABLE = 1 << 0;
        /// Enables monitoring of the coprocessor, typical for x87 instructions.
        ///
        /// Controls together with the `TASK_SWITCHED` flag whether a `wait` or `fwait`
        /// instruction should cause a device-not-available exception.
        const MONITOR_COPROCESSOR = 1 << 1;
        /// Force all x87 and MMX instructions to cause an exception.
        const EMULATE_COPROCESSOR = 1 << 2;
        /// Automatically set to 1 on _hardware_ task switch.
        ///
        /// This flags allows lazily saving x87/MMX/SSE instructions on hardware context switches.
        const TASK_SWITCHED = 1 << 3;
        /// Enables the native error reporting mechanism for x87 FPU errors.
        const NUMERIC_ERROR = 1 << 5;
        /// Controls whether supervisor-level writes to read-only pages are inhibited.
        ///
        /// When set, it is not possible to write to read-only pages from ring 0.
        const WRITE_PROTECT = 1 << 16;
        /// Enables automatic alignment checking.
        const ALIGNMENT_MASK = 1 << 18;
        /// Ignored. Used to control write-back/write-through cache strategy on older CPUs.
        const NOT_WRITE_THROUGH = 1 << 29;
        /// Disables internal caches (only for some cases).
        const CACHE_DISABLE = 1 << 30;
        /// Enables page translation.
        const PAGING = 1 << 31;
    }
}

/// Contains the Page Fault Linear Address (PFLA).
///
/// When page fault occurs, the CPU sets this register to the accessed address.
#[derive(Debug)]
pub struct Cr2;

/// Contains the physical address of the level 4 page table.
#[derive(Debug)]
pub struct Cr3;

bitflags! {
    /// Controls cache settings for the level 4 page table.
    pub struct Cr3Flags: u64 {
        /// Use a writethrough cache policy for the P4 table (else a writeback policy is used).
        const PAGE_LEVEL_WRITETHROUGH = 1 << 3;
        /// Disable caching for the P4 table.
        const PAGE_LEVEL_CACHE_DISABLE = 1 << 4;
    }
}

#[cfg(target_arch = "x86_64")]
mod x86_64 {
    use super::*;
    use crate::structures::paging::PhysFrame;
    use crate::{PhysAddr, VirtAddr};

    impl Cr0 {
        /// Read the current set of CR0 flags.
        pub fn read() -> Cr0Flags {
            Cr0Flags::from_bits_truncate(Self::read_raw())
        }

        /// Read the current raw CR0 value.
        pub fn read_raw() -> u64 {
            let value: u64;
            unsafe {
                asm!("mov %cr0, $0" : "=r" (value));
            }
            value
        }

        /// Write CR0 flags.
        ///
        /// Preserves the value of reserved fields. Unsafe because it's possible to violate memory
        /// safety by e.g. disabling paging.
        pub unsafe fn write(flags: Cr0Flags) {
            let old_value = Self::read_raw();
            let reserved = old_value & !(Cr0Flags::all().bits());
            let new_value = reserved | flags.bits();

            Self::write_raw(new_value);
        }

        /// Write raw CR0 flags.
        ///
        /// Does _not_ preserve any values, including reserved fields. Unsafe because it's possible to violate memory
        /// safety by e.g. disabling paging.
        pub unsafe fn write_raw(value: u64) {
            asm!("mov $0, %cr0" :: "r" (value) : "memory")
        }

        /// Updates CR0 flags.
        ///
        /// Preserves the value of reserved fields. Unsafe because it's possible to violate memory
        /// safety by e.g. disabling paging.
        pub unsafe fn update<F>(f: F)
        where
            F: FnOnce(&mut Cr0Flags),
        {
            let mut flags = Self::read();
            f(&mut flags);
            Self::write(flags);
        }
    }

    impl Cr2 {
        /// Read the current page fault linear address from the CR3 register.
        pub fn read() -> VirtAddr {
            let value: u64;
            unsafe {
                asm!("mov %cr2, $0" : "=r" (value));
            }
            VirtAddr::new(value)
        }
    }

    impl Cr3 {
        /// Read the current P4 table address from the CR3 register.
        pub fn read() -> (PhysFrame, Cr3Flags) {
            let value: u64;
            unsafe {
                asm!("mov %cr3, $0" : "=r" (value));
            }
            let flags = Cr3Flags::from_bits_truncate(value);
            let addr = PhysAddr::new(value & 0x_000f_ffff_ffff_f000);
            let frame = PhysFrame::containing_address(addr);
            (frame, flags)
        }

        /// Write a new P4 table address into the CR3 register.
        ///
        /// ## Safety
        /// Changing the level 4 page table is unsafe, because it's possible to violate memory safety by
        /// changing the page mapping.
        pub unsafe fn write(frame: PhysFrame, flags: Cr3Flags) {
            let addr = frame.start_address();
            let value = addr.as_u64() | flags.bits();
            asm!("mov $0, %cr3" :: "r" (value) : "memory")
        }
    }
}
<commit_msg>Add support for cr4 control register<commit_after>//! Functions to read and write control registers.

pub use super::model_specific::{Efer, EferFlags};

use bitflags::bitflags;

/// Various control flags modifying the basic operation of the CPU.
#[derive(Debug)]
pub struct Cr0;

bitflags! {
    /// Configuration flags of the Cr0 register.
    pub struct Cr0Flags: u64 {
        /// Enables protected mode.
        const PROTECTED_MODE_ENABLE = 1 << 0;
        /// Enables monitoring of the coprocessor, typical for x87 instructions.
        ///
        /// Controls together with the `TASK_SWITCHED` flag whether a `wait` or `fwait`
        /// instruction should cause a device-not-available exception.
        const MONITOR_COPROCESSOR = 1 << 1;
        /// Force all x87 and MMX instructions to cause an exception.
        const EMULATE_COPROCESSOR = 1 << 2;
        /// Automatically set to 1 on _hardware_ task switch.
        ///
        /// This flags allows lazily saving x87/MMX/SSE instructions on hardware context switches.
        const TASK_SWITCHED = 1 << 3;
        /// Enables the native error reporting mechanism for x87 FPU errors.
        const NUMERIC_ERROR = 1 << 5;
        /// Controls whether supervisor-level writes to read-only pages are inhibited.
        ///
        /// When set, it is not possible to write to read-only pages from ring 0.
        const WRITE_PROTECT = 1 << 16;
        /// Enables automatic alignment checking.
        const ALIGNMENT_MASK = 1 << 18;
        /// Ignored. Used to control write-back/write-through cache strategy on older CPUs.
        const NOT_WRITE_THROUGH = 1 << 29;
        /// Disables internal caches (only for some cases).
        const CACHE_DISABLE = 1 << 30;
        /// Enables page translation.
        const PAGING = 1 << 31;
    }
}

/// Contains the Page Fault Linear Address (PFLA).
///
/// When page fault occurs, the CPU sets this register to the accessed address.
#[derive(Debug)]
pub struct Cr2;

/// Contains the physical address of the level 4 page table.
#[derive(Debug)]
pub struct Cr3;

bitflags! {
    /// Controls cache settings for the level 4 page table.
    pub struct Cr3Flags: u64 {
        /// Use a writethrough cache policy for the P4 table (else a writeback policy is used).
        const PAGE_LEVEL_WRITETHROUGH = 1 << 3;
        /// Disable caching for the P4 table.
        const PAGE_LEVEL_CACHE_DISABLE = 1 << 4;
    }
}

/// Various control flags modifying the basic operation of the CPU while in protected mode.
#[derive(Debug)]
pub struct Cr4;

bitflags! {
    /// Controls cache settings for the level 4 page table.
    pub struct Cr4Flags: u64 {
        const VIRTUAL_8086_MODE_EXTENSIONS = 1 << 0;
        const PROTECTED_MODE_VIRTUAL_INTERRUPTS = 1 << 1;
        const TIMESTAMP_DISABLE = 1 << 2;
        const DEBUGGING_EXTENSIONS = 1 << 3;
        const PAGE_SIZE_EXTENSION = 1 << 4;
        const PHYSICAL_ADDRESS_EXTENSION = 1 << 5;
        const MACHINE_CHECK_EXCEPTION = 1 << 6;
        const PAGE_GLOBAL = 1 << 7;
        const PERFORMANCE_MONITOR_COUNTER = 1 << 8;
        const OSFXSR = 1 << 9;
        const OSXMMEXCPT_ENABLE = 1 << 10;
        const USER_MODE_INSTRUCTION_PREVENTION = 1 << 11;
        const L5_PAGING = 1 << 12;
        const VIRTUAL_MACHINE_EXTENSIONS = 1 << 13;
        const SAFER_MODE_EXTENSIONS = 1 << 14;
        const FSGSBASE = 1 << 16;
        const PCID = 1 << 17;
        const OSXSAVE = 1 << 18;
        const SUPERVISOR_MODE_EXECUTION_PROTECTION = 1 << 20;
        const SUPERVISOR_MODE_ACCESS_PREVENTION = 1 << 21;
        const PROTECTION_KEY = 1 << 22;
    }
}

#[cfg(target_arch = "x86_64")]
mod x86_64 {
    use super::*;
    use crate::structures::paging::PhysFrame;
    use crate::{PhysAddr, VirtAddr};

    impl Cr0 {
        /// Read the current set of CR0 flags.
        pub fn read() -> Cr0Flags {
            Cr0Flags::from_bits_truncate(Self::read_raw())
        }

        /// Read the current raw CR0 value.
        pub fn read_raw() -> u64 {
            let value: u64;
            unsafe {
                asm!("mov %cr0, $0" : "=r" (value));
            }
            value
        }

        /// Write CR0 flags.
        ///
        /// Preserves the value of reserved fields. Unsafe because it's possible to violate memory
        /// safety by e.g. disabling paging.
        pub unsafe fn write(flags: Cr0Flags) {
            let old_value = Self::read_raw();
            let reserved = old_value & !(Cr0Flags::all().bits());
            let new_value = reserved | flags.bits();

            Self::write_raw(new_value);
        }

        /// Write raw CR0 flags.
        ///
        /// Does _not_ preserve any values, including reserved fields. Unsafe because it's possible to violate memory
        /// safety by e.g. disabling paging.
        pub unsafe fn write_raw(value: u64) {
            asm!("mov $0, %cr0" :: "r" (value) : "memory")
        }

        /// Updates CR0 flags.
        ///
        /// Preserves the value of reserved fields. Unsafe because it's possible to violate memory
        /// safety by e.g. disabling paging.
        pub unsafe fn update<F>(f: F)
        where
            F: FnOnce(&mut Cr0Flags),
        {
            let mut flags = Self::read();
            f(&mut flags);
            Self::write(flags);
        }
    }

    impl Cr2 {
        /// Read the current page fault linear address from the CR3 register.
        pub fn read() -> VirtAddr {
            let value: u64;
            unsafe {
                asm!("mov %cr2, $0" : "=r" (value));
            }
            VirtAddr::new(value)
        }
    }

    impl Cr3 {
        /// Read the current P4 table address from the CR3 register.
        pub fn read() -> (PhysFrame, Cr3Flags) {
            let value: u64;
            unsafe {
                asm!("mov %cr3, $0" : "=r" (value));
            }
            let flags = Cr3Flags::from_bits_truncate(value);
            let addr = PhysAddr::new(value & 0x_000f_ffff_ffff_f000);
            let frame = PhysFrame::containing_address(addr);
            (frame, flags)
        }

        /// Write a new P4 table address into the CR3 register.
        ///
        /// ## Safety
        /// Changing the level 4 page table is unsafe, because it's possible to violate memory safety by
        /// changing the page mapping.
        pub unsafe fn write(frame: PhysFrame, flags: Cr3Flags) {
            let addr = frame.start_address();
            let value = addr.as_u64() | flags.bits();
            asm!("mov $0, %cr3" :: "r" (value) : "memory")
        }
    }

    impl Cr4 {
        /// Read the current set of CR4 flags.
        pub fn read() -> Cr4Flags {
            Cr4Flags::from_bits_truncate(Self::read_raw())
        }

        /// Read the current raw CR4 value.
        pub fn read_raw() -> u64 {
            let value: u64;
            unsafe {
                asm!("mov %cr4, $0" : "=r" (value));
            }
            value
        }

        /// Write CR4 flags.
        ///
        /// Preserves the value of reserved fields. Unsafe because it's possible to violate memory
        /// safety by e.g. physical address extension.
        pub unsafe fn write(flags: Cr4Flags) {
            let old_value = Self::read_raw();
            let reserved = old_value & !(Cr4Flags::all().bits());
            let new_value = reserved | flags.bits();

            Self::write_raw(new_value);
        }

        /// Write raw CR4 flags.
        ///
        /// Does _not_ preserve any values, including reserved fields. Unsafe because it's possible to violate memory
        /// safety by e.g. physical address extension.
        pub unsafe fn write_raw(value: u64) {
            asm!("mov $0, %cr4" :: "r" (value) : "memory")
        }

        /// Updates CR4 flags.
        ///
        /// Preserves the value of reserved fields. Unsafe because it's possible to violate memory
        /// safety by e.g. physical address extension.
        pub unsafe fn update<F>(f: F)
        where
            F: FnOnce(&mut Cr4Flags),
        {
            let mut flags = Self::read();
            f(&mut flags);
            Self::write(flags);
        }
    }
}
<|endoftext|>


<commit_before>use std::sync::Future;
use std::fmt;

use numerics;
use parse::IrcMsg;
use message_types::server;
use watchers::base::EventWatcher;
use event::IrcEvent;

pub type RegisterResult = Result<(), RegisterError>;

#[deriving(Clone, Send, Show)]
pub struct RegisterError {
    pub errtype: RegisterErrorType,
    pub message: IrcMsg,
}

impl RegisterError {
    pub fn should_pick_new_nickname(&self) -> bool {
        match server::IncomingMsg::from_msg(self.message.clone()) {
            server::IncomingMsg::Numeric(num, _) => {
                numerics::ERR_NICKNAMEINUSE == (num as i32)
            },
            _ => false
        }
    }
}

#[deriving(Clone, Send, Show, Copy)]
pub enum RegisterErrorType {
    NoNicknameGiven,
    NicknameInUse,
    UnavailableResource,
    ErroneousNickname,
    NicknameCollision,
    Restricted,
}

impl RegisterErrorType {

    pub fn is_known_error(result: i32) -> bool {
        RegisterErrorType::from_ord_known(result).is_some()
    }

    pub fn from_ord_known(result: i32) -> Option<RegisterErrorType> {
        match result {
            numerics::ERR_NONICKNAMEGIVEN => Some(RegisterErrorType::NoNicknameGiven),
            numerics::ERR_NICKNAMEINUSE => Some(RegisterErrorType::NicknameInUse),
            numerics::ERR_UNAVAILRESOURCE => Some(RegisterErrorType::UnavailableResource),
            numerics::ERR_ERRONEUSNICKNAME => Some(RegisterErrorType::ErroneousNickname),
            numerics::ERR_NICKCOLLISION => Some(RegisterErrorType::NicknameCollision),
            numerics::ERR_RESTRICTED => Some(RegisterErrorType::Restricted),
            _ => None
        }
    }
}

pub struct RegisterEventWatcher {
    rx_connect: bool,
    result: Option<Result<(), RegisterError>>,
    monitors: Vec<SyncSender<Result<(), RegisterError>>>
}

impl RegisterEventWatcher {
    pub fn new() -> RegisterEventWatcher {
        RegisterEventWatcher {
            rx_connect: false,
            result: None,
            monitors: Vec::new()
        }
    }

    fn dispatch_monitors(&mut self) {
        let result = self.result.clone().unwrap();
        for monitor in self.monitors.iter() {
            match monitor.try_send(result.clone()) {
                Ok(_) => (),
                Err(_) => panic!("sending failed")
            }
        }
        self.monitors = Vec::new();
    }

    fn add_monitor(&mut self, monitor: SyncSender<Result<(), RegisterError>>) {
        let result = self.result.clone();

        match (self.rx_connect, result) {
            (true, Some(result)) => {
                monitor.send(result.clone());
            },
            (true, None) => {
                panic!("rx_connect without result");
            },
            (false, Some(_)) => {
                panic!("result without rx_connect");
            },
            (false, None) => {
                self.monitors.push(monitor);
            }
        };
    }

    pub fn get_monitor(&mut self) -> Receiver<RegisterResult> {
        let (tx, rx) = sync_channel(1);
        self.add_monitor(tx);
        rx
    }

    pub fn get_future(&mut self) -> Future<RegisterResult> {
        Future::from_receiver(self.get_monitor())
    }

    fn accept_ircmessage(&mut self, msg: &IrcMsg) {
        println!("RegisterEventWatcher: RX {}", msg);

        let (interested, err) = match server::IncomingMsg::from_msg(msg.clone()) {
            server::IncomingMsg::Numeric(1, _) => {
                (true, None)
            },
            server::IncomingMsg::Numeric(other, _) => {
                let res = RegisterErrorType::from_ord_known(other as i32);
                (res.is_some(), res)
            },
            _ => (false, None)
        };

        if interested {
            self.rx_connect = true;
            self.result = Some(match err {
                None => Ok(()),
                Some(err) => Err(RegisterError {
                    errtype: err,
                    message: msg.clone()
                })
            });
            self.dispatch_monitors();
        }
    }
}

impl EventWatcher for RegisterEventWatcher {
    fn on_event(&mut self, event: &IrcEvent) {
        match *event {
            IrcEvent::IrcMsg(ref msg) => {
                self.accept_ircmessage(msg);
            },
            _ => ()
        }
    }

    fn is_finished(&self) -> bool {
       self.rx_connect
    }

    fn get_name(&self) -> &'static str {
        "RegisterEventWatcher"
    }

    fn display(&self) -> String {
        format!("{}", self)
    }
}

impl fmt::Show for RegisterEventWatcher {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "RegisterEventWatcher()")
    }
}<commit_msg>Remove deriving(Send) from registration structs<commit_after>use std::sync::Future;
use std::fmt;

use numerics;
use parse::IrcMsg;
use message_types::server;
use watchers::base::EventWatcher;
use event::IrcEvent;

pub type RegisterResult = Result<(), RegisterError>;

#[deriving(Clone, Show)]
pub struct RegisterError {
    pub errtype: RegisterErrorType,
    pub message: IrcMsg,
}

impl RegisterError {
    pub fn should_pick_new_nickname(&self) -> bool {
        match server::IncomingMsg::from_msg(self.message.clone()) {
            server::IncomingMsg::Numeric(num, _) => {
                numerics::ERR_NICKNAMEINUSE == (num as i32)
            },
            _ => false
        }
    }
}

#[deriving(Clone, Show, Copy)]
pub enum RegisterErrorType {
    NoNicknameGiven,
    NicknameInUse,
    UnavailableResource,
    ErroneousNickname,
    NicknameCollision,
    Restricted,
}

impl RegisterErrorType {

    pub fn is_known_error(result: i32) -> bool {
        RegisterErrorType::from_ord_known(result).is_some()
    }

    pub fn from_ord_known(result: i32) -> Option<RegisterErrorType> {
        match result {
            numerics::ERR_NONICKNAMEGIVEN => Some(RegisterErrorType::NoNicknameGiven),
            numerics::ERR_NICKNAMEINUSE => Some(RegisterErrorType::NicknameInUse),
            numerics::ERR_UNAVAILRESOURCE => Some(RegisterErrorType::UnavailableResource),
            numerics::ERR_ERRONEUSNICKNAME => Some(RegisterErrorType::ErroneousNickname),
            numerics::ERR_NICKCOLLISION => Some(RegisterErrorType::NicknameCollision),
            numerics::ERR_RESTRICTED => Some(RegisterErrorType::Restricted),
            _ => None
        }
    }
}

pub struct RegisterEventWatcher {
    rx_connect: bool,
    result: Option<Result<(), RegisterError>>,
    monitors: Vec<SyncSender<Result<(), RegisterError>>>
}

impl RegisterEventWatcher {
    pub fn new() -> RegisterEventWatcher {
        RegisterEventWatcher {
            rx_connect: false,
            result: None,
            monitors: Vec::new()
        }
    }

    fn dispatch_monitors(&mut self) {
        let result = self.result.clone().unwrap();
        for monitor in self.monitors.iter() {
            match monitor.try_send(result.clone()) {
                Ok(_) => (),
                Err(_) => panic!("sending failed")
            }
        }
        self.monitors = Vec::new();
    }

    fn add_monitor(&mut self, monitor: SyncSender<Result<(), RegisterError>>) {
        let result = self.result.clone();

        match (self.rx_connect, result) {
            (true, Some(result)) => {
                monitor.send(result.clone());
            },
            (true, None) => {
                panic!("rx_connect without result");
            },
            (false, Some(_)) => {
                panic!("result without rx_connect");
            },
            (false, None) => {
                self.monitors.push(monitor);
            }
        };
    }

    pub fn get_monitor(&mut self) -> Receiver<RegisterResult> {
        let (tx, rx) = sync_channel(1);
        self.add_monitor(tx);
        rx
    }

    pub fn get_future(&mut self) -> Future<RegisterResult> {
        Future::from_receiver(self.get_monitor())
    }

    fn accept_ircmessage(&mut self, msg: &IrcMsg) {
        println!("RegisterEventWatcher: RX {}", msg);

        let (interested, err) = match server::IncomingMsg::from_msg(msg.clone()) {
            server::IncomingMsg::Numeric(1, _) => {
                (true, None)
            },
            server::IncomingMsg::Numeric(other, _) => {
                let res = RegisterErrorType::from_ord_known(other as i32);
                (res.is_some(), res)
            },
            _ => (false, None)
        };

        if interested {
            self.rx_connect = true;
            self.result = Some(match err {
                None => Ok(()),
                Some(err) => Err(RegisterError {
                    errtype: err,
                    message: msg.clone()
                })
            });
            self.dispatch_monitors();
        }
    }
}

impl EventWatcher for RegisterEventWatcher {
    fn on_event(&mut self, event: &IrcEvent) {
        match *event {
            IrcEvent::IrcMsg(ref msg) => {
                self.accept_ircmessage(msg);
            },
            _ => ()
        }
    }

    fn is_finished(&self) -> bool {
       self.rx_connect
    }

    fn get_name(&self) -> &'static str {
        "RegisterEventWatcher"
    }

    fn display(&self) -> String {
        format!("{}", self)
    }
}

impl fmt::Show for RegisterEventWatcher {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "RegisterEventWatcher()")
    }
}
<|endoftext|>