<commit_before>use rustc::lint::{EarlyContext, EarlyLintPass, LintArray, LintContext, LintPass};
use std::cmp::Ordering;
use syntax::ast::{Item, Item_, Lit_, MetaItem_, Mod, PathListItem_, ViewPath_, Visibility};
use syntax::codemap::Span;
use syntax::print::pprust::path_to_string;

// Warn unsorted declarations by default (since denying is a poor choice for styling lints)
declare_lint!(UNSORTED_DECLARATIONS, Warn,
              "Warn when the declarations of crates or modules are not in alphabetical order");

pub struct Sorty;

impl LintPass for Sorty {
    fn get_lints(&self) -> LintArray {
        lint_array!(UNSORTED_DECLARATIONS)
    }
}

impl EarlyLintPass for Sorty {
    // walking through all the modules is enough for our purpose
    fn check_mod(&mut self, cx: &EarlyContext, module: &Mod, _span: Span, _id: u32) {
        // TODO: lint should stop ignoring the comments near the declarations
        let session_codemap = cx.sess.codemap();    // required only for checking inline mods
        let mut extern_crates = Vec::new();
        let mut uses = Vec::new();
        let mut mods = Vec::new();

        for item in &module.items {
            // I've used `format!` most of the time, because we have a mixture of Strings & InternedStrings
            let item_name = format!("{}", item.ident.name.as_str());
            let item_span = item.span;
            match item.node.clone() {
                Item_::ItemExternCrate(optional_name) if item_name != "std" => {
                    // We've put the declaration here because, we have to sort crate declarations
                    // with respect to the renamed version (instead of the old one).
                    // Since we also don't have `pub` (indicated by the `false` below),
                    // we could just append the declaration to the attributes.
                    let mut item_attrs = get_item_attrs(&item, false);
                    item_attrs = match optional_name {    // for `extern crate foo as bar`
                        Some(old_name) => format!("{}extern crate {} as", item_attrs, old_name),
                        None => format!("{}extern crate", item_attrs),
                    }; extern_crates.push((item_name, item_attrs, item_span, false));
                },

                Item_::ItemMod(module) => {
                    let mod_invoked_file = session_codemap.span_to_filename(item.span);
                    let mod_declared_file = session_codemap.span_to_filename(module.inner);
                    if mod_declared_file != mod_invoked_file {      // ignores inline modules
                        let item_attrs = get_item_attrs(&item, true);
                        mods.push((item_name, item_attrs, item_span, false));
                    }
                },

                Item_::ItemUse(spanned) => {
                    let item_attrs = get_item_attrs(&item, true);
                    match spanned.node {
                        ViewPath_::ViewPathSimple(ref ident, ref path) => {
                            let path_str = path_to_string(&path);
                            let name = ident.name.as_str();
                            let renamed = {     // for `use foo as bar`
                                let split = path_str.split(":").collect::<Vec<&str>>();
                                match split[split.len() - 1] == &*name {
                                    true => path_str.clone(),
                                    false => format!("{} as {}", &path_str, &name),
                                }
                            }; uses.push((renamed, item_attrs, item_span, false));
                        },

                        ViewPath_::ViewPathList(ref path, ref list) => {
                            // TODO: track the renamed items in use lists
                            let old_list = list
                                           .iter()
                                           .map(|&list_item| {
                                                match list_item.node {
                                                    PathListItem_::PathListMod { .. } =>
                                                        "self".to_owned(),      // this must be `self`
                                                    PathListItem_::PathListIdent { name, .. } => {
                                                        let interned = name.name.as_str();
                                                        let string = &*interned;
                                                        string.to_owned()
                                                    },
                                                }
                                            }).collect::<Vec<String>>();

                            let mut new_list = old_list.clone();
                            new_list.sort_by(|a, b| {
                                match (&**a, &**b) {    // `self` should be first in an use list
                                    ("self", _) => Ordering::Less,
                                    (_, "self") => Ordering::Greater,
                                    _ => a.cmp(b),
                                }
                            });

                            let mut warn = false;
                            let use_list = format!("{}::{{{}}}", path_to_string(&path), new_list.join(", "));
                            for i in 0..old_list.len() {
                                if old_list[i] != new_list[i] {
                                    warn = true;    // check whether the use list is sorted
                                    break;
                                }
                            } uses.push((use_list, item_attrs, path.span, warn));
                        },

                        ViewPath_::ViewPathGlob(ref path) => {
                            let path_str = path_to_string(&path) + "::*";
                            // we don't have any use statements like `use std::prelude::*`
                            // since it's done only by rustc, we can safely neglect those here
                            if !path_str.starts_with("std::") {
                                uses.push((path_str, item_attrs, item_span, false));
                            }
                        },
                    }
                },
                _ => (),
            }
        }

        // we don't include the crate declaration here, because we've already appended it with the attributes
        check_sort(&extern_crates, cx, "crate declarations", "");
        check_sort(&mods, cx, "module declarations (other than inline modules)", "mod");
        check_sort(&uses, cx, "use statements", "use");

        // for collecting, formatting & filtering the attributes (and checking the visibility)
        fn get_item_attrs(item: &Item, pub_check: bool) -> String {
            let mut attr_vec = item.attrs
                               .iter()
                               .filter_map(|attr| {
                                   let meta_item = attr.node.value.node.clone();
                                   let meta_string = get_meta_as_string(&meta_item);
                                   match meta_string.starts_with("doc = ") {
                                       true => None,
                                       false => Some(format!("#[{}]", meta_string)),
                                   }
                               }).collect::<Vec<String>>();

            attr_vec.sort_by(|a, b| {
                match (&**a, &**b) {    // put `macro_use` first for later checking
                    ("#[macro_use]", _) => Ordering::Less,
                    (_, "#[macro_use]") => Ordering::Greater,
                    _ => a.cmp(b),
                }
            }); let attr_string = attr_vec.join("\n");

            match item.vis {
                Visibility::Public if pub_check => {
                    match attr_string.is_empty() {
                        true => "pub ".to_owned(),
                        false => attr_string + "\npub ",    // `pub` for mods and uses
                    }
                },
                _ => {
                    match attr_string.is_empty() {
                        true => attr_string,
                        false => attr_string + "\n",
                    }
                },
            }
        }

        // collect the information from meta items into Strings
        fn get_meta_as_string(meta_item: &MetaItem_) -> String {
            match *meta_item {
                MetaItem_::MetaWord(ref string) => format!("{}", string),
                MetaItem_::MetaList(ref string, ref meta_items) => {
                    let stuff = meta_items
                                .iter()
                                .map(|meta_item| {
                                    get_meta_as_string(&meta_item.node)
                                }).collect::<Vec<String>>();
                    format!("{}({})", string, stuff.join(", "))
                },
                MetaItem_::MetaNameValue(ref string, ref literal) => {
                    let value = match literal.node {
                        Lit_::LitStr(ref inner_str, _style) => inner_str,
                        _ => panic!("unexpected literal found for meta item!"),     // which doesn't really happen
                    }; format!("{} = \"{}\"", string, value)
                },
            }
        }

        // checks the sorting of all the declarations and raises warnings whenever necessary
        // takes a slice of tuples with name, related attributes, spans and whether to warn for unordered use lists
        fn check_sort(old_list: &[(String, String, Span, bool)], cx: &EarlyContext, kind: &str, syntax: &str) {
            let length = old_list.len();
            let mut new_list = old_list
                                .iter()
                                .map(|&(ref name, ref attrs, _span, warn)| (name.clone(), attrs.clone(), warn))
                                .collect::<Vec<(String, String, bool)>>();
            new_list.sort_by(|&(ref str_a, ref attr_a, _), &(ref str_b, ref attr_b, _)| {
                // move the `pub` statements below (with `~` since it's on the farther side of ASCII)
                let mut new_str_a = str_for_biased_sort(&str_a, attr_a.ends_with("pub "), "~");
                let mut new_str_b = str_for_biased_sort(&str_b, attr_b.ends_with("pub "), "~");
                // move the #[macro_use] stuff above (with `!` since it's on the lower extreme of ASCII)
                new_str_a = str_for_biased_sort(&new_str_a, attr_a.starts_with("#[macro_use]"), "!");
                new_str_b = str_for_biased_sort(&new_str_b, attr_b.starts_with("#[macro_use]"), "!");
                new_str_a.cmp(&new_str_b)
            });

            let mut index = 0;
            let mut span: Option<Span> = None;
            for i in 0..length {
                if (old_list[i].0 != new_list[i].0) || new_list[i].2 {
                    span = Some(old_list[i].2);
                    index = i;      // only to find the index of the first unsorted declaration
                    break;          // because, we'll be printing everything following the first unsorted one
                }
            }

            match span {
                Some(span_start) => {   // print all the declarations proceeding the first unsorted one
                    let suggestion_list = (index..length)
                                          .map(|i| {
                                              if i == length - 1 {  // increase the span to include more lines
                                                  let mut sp = span_start;
                                                  sp.hi = old_list[i].2.hi;
                                                  span = Some(sp);
                                              } format!("{}{} {};", new_list[i].1, syntax, new_list[i].0)
                                          }).collect::<Vec<String>>();
                    let message = format!("{} should be in alphabetical order!", kind);
                    let suggestion = format!("Try this...\n\n{}\n", suggestion_list.join("\n"));
                    // unwrapping the value here, because it's quite certain that there's something in `span`
                    cx.span_lint_help(UNSORTED_DECLARATIONS, span_start, &message, &suggestion);
                },
                None => (),
            }

            // prepend given characters to the names for "biased" sorting
            fn str_for_biased_sort(string: &String, choice: bool, prepend_char: &str) -> String {
                match choice {
                    true => prepend_char.to_owned() + &string,
                    false => string.clone()
                }
            }
        }
    }
}
<commit_msg>run rustfmt<commit_after>use rustc::lint::{EarlyContext, EarlyLintPass, LintArray, LintContext, LintPass};
use std::cmp::Ordering;
use syntax::ast::{Item, Item_, Lit_, MetaItem_, Mod, PathListItem_, ViewPath_, Visibility};
use syntax::codemap::Span;
use syntax::print::pprust::path_to_string;

// Warn unsorted declarations by default (since denying is a poor choice for styling lints)
declare_lint!(UNSORTED_DECLARATIONS,
              Warn,
              "Warn when the declarations of crates or modules are not in alphabetical order");

pub struct Sorty;

impl LintPass for Sorty {
    fn get_lints(&self) -> LintArray {
        lint_array!(UNSORTED_DECLARATIONS)
    }
}

impl EarlyLintPass for Sorty {
    // walking through all the modules is enough for our purpose
    fn check_mod(&mut self, cx: &EarlyContext, module: &Mod, _span: Span, _id: u32) {
        // TODO: lint should stop ignoring the comments near the declarations
        let session_codemap = cx.sess.codemap();    // required only for checking inline mods
        let mut extern_crates = Vec::new();
        let mut uses = Vec::new();
        let mut mods = Vec::new();

        for item in &module.items {
            // I've used `format!` most of the time, because we have a mixture of Strings &
            // InternedStrings
            let item_name = format!("{}", item.ident.name.as_str());
            let item_span = item.span;
            match item.node.clone() {
                Item_::ItemExternCrate(optional_name) if item_name != "std" => {
                    // We've put the declaration here because, we have to sort crate declarations
                    // with respect to the renamed version (instead of the old one).
                    // Since we also don't have `pub` (indicated by the `false` below),
                    // we could just append the declaration to the attributes.
                    let mut item_attrs = get_item_attrs(&item, false);
                    item_attrs = match optional_name {    // for `extern crate foo as bar`
                        Some(old_name) => format!("{}extern crate {} as", item_attrs, old_name),
                        None => format!("{}extern crate", item_attrs),
                    };
                    extern_crates.push((item_name, item_attrs, item_span, false));
                }

                Item_::ItemMod(module) => {
                    let mod_invoked_file = session_codemap.span_to_filename(item.span);
                    let mod_declared_file = session_codemap.span_to_filename(module.inner);
                    if mod_declared_file != mod_invoked_file {
                        // ignores inline modules
                        let item_attrs = get_item_attrs(&item, true);
                        mods.push((item_name, item_attrs, item_span, false));
                    }
                }

                Item_::ItemUse(spanned) => {
                    let item_attrs = get_item_attrs(&item, true);
                    match spanned.node {
                        ViewPath_::ViewPathSimple(ref ident, ref path) => {
                            let path_str = path_to_string(&path);
                            let name = ident.name.as_str();
                            let renamed = {
                                // for `use foo as bar`
                                let split = path_str.split(":").collect::<Vec<&str>>();
                                match split[split.len() - 1] == &*name {
                                    true => path_str.clone(),
                                    false => format!("{} as {}", &path_str, &name),
                                }
                            };
                            uses.push((renamed, item_attrs, item_span, false));
                        }

                        ViewPath_::ViewPathList(ref path, ref list) => {
                            // TODO: track the renamed items in use lists
                            let old_list =
                                list.iter()
                                    .map(|&list_item| {
                                        match list_item.node {
                                            PathListItem_::PathListMod { .. } => {
                                                // this must be `self`
                                                "self".to_owned()
                                            }
                                            PathListItem_::PathListIdent { name, .. } => {
                                                let interned = name.name.as_str();
                                                let string = &*interned;
                                                string.to_owned()
                                            }
                                        }
                                    })
                                    .collect::<Vec<String>>();

                            let mut new_list = old_list.clone();
                            new_list.sort_by(|a, b| {
                                match (&**a, &**b) {    // `self` should be first in an use list
                                    ("self", _) => Ordering::Less,
                                    (_, "self") => Ordering::Greater,
                                    _ => a.cmp(b),
                                }
                            });

                            let mut warn = false;
                            let use_list = format!("{}::{{{}}}",
                                                   path_to_string(&path),
                                                   new_list.join(", "));
                            for i in 0..old_list.len() {
                                if old_list[i] != new_list[i] {
                                    warn = true;    // check whether the use list is sorted
                                    break;
                                }
                            }
                            uses.push((use_list, item_attrs, path.span, warn));
                        }

                        ViewPath_::ViewPathGlob(ref path) => {
                            let path_str = path_to_string(&path) + "::*";
                            // we don't have any use statements like `use std::prelude::*`
                            // since it's done only by rustc, we can safely neglect those here
                            if !path_str.starts_with("std::") {
                                uses.push((path_str, item_attrs, item_span, false));
                            }
                        }
                    }
                }
                _ => (),
            }
        }

        // we don't include the crate declaration here, because we've already appended it with the
        // attributes
        check_sort(&extern_crates, cx, "crate declarations", "");
        check_sort(&mods,
                   cx,
                   "module declarations (other than inline modules)",
                   "mod");
        check_sort(&uses, cx, "use statements", "use");

        // for collecting, formatting & filtering the attributes (and checking the visibility)
        fn get_item_attrs(item: &Item, pub_check: bool) -> String {
            let mut attr_vec = item.attrs
                                   .iter()
                                   .filter_map(|attr| {
                                       let meta_item = attr.node.value.node.clone();
                                       let meta_string = get_meta_as_string(&meta_item);
                                       match meta_string.starts_with("doc = ") {
                                           true => None,
                                           false => Some(format!("#[{}]", meta_string)),
                                       }
                                   })
                                   .collect::<Vec<String>>();

            attr_vec.sort_by(|a, b| {
                match (&**a, &**b) {    // put `macro_use` first for later checking
                    ("#[macro_use]", _) => Ordering::Less,
                    (_, "#[macro_use]") => Ordering::Greater,
                    _ => a.cmp(b),
                }
            });
            let attr_string = attr_vec.join("\n");

            match item.vis {
                Visibility::Public if pub_check => {
                    match attr_string.is_empty() {
                        true => "pub ".to_owned(),
                        false => attr_string + "\npub ",    // `pub` for mods and uses
                    }
                }
                _ => {
                    match attr_string.is_empty() {
                        true => attr_string,
                        false => attr_string + "\n",
                    }
                }
            }
        }

        // collect the information from meta items into Strings
        fn get_meta_as_string(meta_item: &MetaItem_) -> String {
            match *meta_item {
                MetaItem_::MetaWord(ref string) => format!("{}", string),
                MetaItem_::MetaList(ref string, ref meta_items) => {
                    let stuff = meta_items.iter()
                                          .map(|meta_item| get_meta_as_string(&meta_item.node))
                                          .collect::<Vec<String>>();
                    format!("{}({})", string, stuff.join(", "))
                }
                MetaItem_::MetaNameValue(ref string, ref literal) => {
                    let value = match literal.node {
                        Lit_::LitStr(ref inner_str, _style) => inner_str,
                        _ => {
                            // doesn't really happen
                            panic!("unexpected literal found for meta item!")
                        }
                    };
                    format!("{} = \"{}\"", string, value)
                }
            }
        }

        // checks the sorting of all the declarations and raises warnings whenever necessary
        // takes a slice of tuples with name, related attributes, spans and whether to warn for
        // unordered use lists
        fn check_sort(old_list: &[(String, String, Span, bool)],
                      cx: &EarlyContext,
                      kind: &str,
                      syntax: &str) {
            let length = old_list.len();
            let mut new_list = old_list.iter()
                                       .map(|&(ref name, ref attrs, _span, warn)| {
                                           (name.clone(), attrs.clone(), warn)
                                       })
                                       .collect::<Vec<(String, String, bool)>>();
            new_list.sort_by(|&(ref str_a, ref attr_a, _), &(ref str_b, ref attr_b, _)| {
                // move the `pub` statements below (with `~` since it's on the farther side of
                // ASCII)
                let mut new_str_a = str_for_biased_sort(&str_a, attr_a.ends_with("pub "), "~");
                let mut new_str_b = str_for_biased_sort(&str_b, attr_b.ends_with("pub "), "~");
                // move the #[macro_use] stuff above (with `!` since it's on the lower extreme of
                // ASCII)
                new_str_a = str_for_biased_sort(&new_str_a,
                                                attr_a.starts_with("#[macro_use]"),
                                                "!");
                new_str_b = str_for_biased_sort(&new_str_b,
                                                attr_b.starts_with("#[macro_use]"),
                                                "!");
                new_str_a.cmp(&new_str_b)
            });

            let mut index = 0;
            let mut span: Option<Span> = None;
            for i in 0..length {
                if (old_list[i].0 != new_list[i].0) || new_list[i].2 {
                    span = Some(old_list[i].2);
                    // only to find the index of the first unsorted declaration
                    index = i;
                    // because we'll be printing everything following the first unsorted one
                    break;
                }
            }

            match span {
                Some(span_start) => {
                    // print all the declarations proceeding the first unsorted one
                    let suggestion_list = (index..length)
                                              .map(|i| {
                                                  if i == length - 1 {
                                                      // increase the span to include more lines
                                                      let mut sp = span_start;
                                                      sp.hi = old_list[i].2.hi;
                                                      span = Some(sp);
                                                  }
                                                  format!("{}{} {};",
                                                          new_list[i].1,
                                                          syntax,
                                                          new_list[i].0)
                                              })
                                              .collect::<Vec<String>>();
                    let message = format!("{} should be in alphabetical order!", kind);
                    let suggestion = format!("Try this...\n\n{}\n", suggestion_list.join("\n"));
                    // unwrapping the value here, because it's quite certain that there's something
                    // in `span`
                    cx.span_lint_help(UNSORTED_DECLARATIONS, span_start, &message, &suggestion);
                }
                None => (),
            }

            // prepend given characters to the names for "biased" sorting
            fn str_for_biased_sort(string: &String, choice: bool, prepend_char: &str) -> String {
                match choice {
                    true => prepend_char.to_owned() + &string,
                    false => string.clone(),
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::option::Option;
use std::sync::{Arc, RwLock};
use std::boxed::Box;
use std::net::SocketAddr;

use mio::{Token, Handler, EventLoop, EventSet, PollOpt};
use mio::tcp::{TcpListener, TcpStream};

use raftserver::store::{cmd_resp, Transport};
use raftserver::Result as RaftResult;
use kvproto::raft_cmdpb::{RaftCmdRequest, RaftCmdResponse};
use kvproto::msgpb::{MessageType, Message};
use super::{Msg, SendCh, ConnData};
use super::conn::Conn;
use super::{Result, other};
use util::HandyRwLock;
use storage::Storage;
use super::kv::StoreHandler;

const SERVER_TOKEN: Token = Token(1);
const FIRST_CUSTOM_TOKEN: Token = Token(1024);
const INVALID_TOKEN: Token = Token(0);

pub fn create_event_loop<T: Transport>() -> Result<EventLoop<Server<T>>> {
    let event_loop = try!(EventLoop::new());
    Ok(event_loop)
}

pub struct Server<T: Transport> {
    listener: TcpListener,
    // We use HashMap instead of common use mio slab to avoid token reusing.
    // In our raft server, a client with token 1 sends a raft command, we will
    // propose this command, execute it then send the response to the client with
    // token 1. But before the response, the client connection is broken and another
    // new client connects, mio slab may reuse the token 1 for it. So the subsequent
    // response will be sent to the new client.
    // To avoid this, we use the HashMap instead and can guarantee the token id is
    // unique and can't be reused.
    conns: HashMap<Token, Conn>,
    conn_token_counter: usize,
    sendch: SendCh,

    // addr -> Token
    peers: HashMap<String, Token>,

    trans: Arc<RwLock<T>>,

    store: StoreHandler,
}

impl<T: Transport> Server<T> {
    // Create a server with already initialized engines.
    // We must bootstrap all stores before running the server.
    pub fn new(event_loop: &mut EventLoop<Self>,
               addr: &str,
               storage: Storage,
               trans: Arc<RwLock<T>>)
               -> Result<Server<T>> {
        let laddr = try!(addr.parse());
        let listener = try!(TcpListener::bind(&laddr));

        try!(event_loop.register(&listener,
                                 SERVER_TOKEN,
                                 EventSet::readable(),
                                 PollOpt::edge()));

        let sendch = SendCh::new(event_loop.channel());

        let store_handler = StoreHandler::new(storage, sendch.clone());

        let svr = Server {
            listener: listener,
            sendch: sendch,
            conns: HashMap::new(),
            conn_token_counter: FIRST_CUSTOM_TOKEN.as_usize(),
            peers: HashMap::new(),
            trans: trans,
            store: store_handler,
        };

        Ok(svr)
    }

    pub fn get_trans(&self) -> Arc<RwLock<T>> {
        self.trans.clone()
    }

    pub fn run(&mut self, event_loop: &mut EventLoop<Self>) -> Result<()> {
        try!(event_loop.run(self));
        Ok(())
    }

    pub fn get_sendch(&self) -> SendCh {
        self.sendch.clone()
    }

    // Return listening address, this may only be used for outer test
    // to get the real address because we may use "127.0.0.1:0"
    // in test to avoid port conflict.
    pub fn listening_addr(&self) -> Result<SocketAddr> {
        let addr = try!(self.listener.local_addr());
        Ok(addr)
    }

    fn remove_conn(&mut self, event_loop: &mut EventLoop<Self>, token: Token) {
        let conn = self.conns.remove(&token);
        match conn {
            Some(conn) => {
                // if connected to remote peer, remove this too.
                if let Some(addr) = conn.peer_addr {
                    self.peers.remove(&addr);
                }

                if let Err(e) = event_loop.deregister(&conn.sock) {
                    error!("deregister conn err {:?}", e);
                }
            }
            None => {
                warn!("missing connection for token {}", token.as_usize());
            }
        }
    }

    fn add_new_conn(&mut self,
                    event_loop: &mut EventLoop<Self>,
                    sock: TcpStream,
                    peer_addr: Option<String>)
                    -> Result<Token> {
        let new_token = Token(self.conn_token_counter);
        self.conn_token_counter += 1;

        // TODO: check conn max capacity.

        try!(sock.set_nodelay(true));

        try!(event_loop.register(&sock,
                                 new_token,
                                 EventSet::readable() | EventSet::hup(),
                                 PollOpt::edge()));

        let conn = Conn::new(sock, new_token, peer_addr);
        self.conns.insert(new_token, conn);

        Ok(new_token)
    }

    fn on_conn_readable(&mut self, event_loop: &mut EventLoop<Self>, token: Token) -> Result<()> {
        let msgs = try!(match self.conns.get_mut(&token) {
            None => {
                warn!("missing conn for token {:?}", token);
                return Ok(());
            }
            Some(conn) => conn.read(event_loop),
        });

        if msgs.is_empty() {
            // Read no message, no need to handle.
            return Ok(());
        }

        for msg in msgs {
            try!(self.on_conn_msg(token, msg))
        }

        Ok(())
    }

    fn on_conn_msg(&mut self, token: Token, data: ConnData) -> Result<()> {
        let msg_id = data.msg_id;
        let mut msg = data.msg;

        let msg_type = msg.get_msg_type();
        match msg_type {
            MessageType::Raft => {
                if let Err(e) = self.trans.rl().send_raft_msg(msg.take_raft()) {
                    // Should we return error to let outer close this connection later?
                    error!("send raft message for token {:?} with msg id {} err {:?}",
                           token,
                           msg_id,
                           e);
                }
                Ok(())
            }
            MessageType::Cmd => self.on_raft_command(msg.take_cmd_req(), token, msg_id),
            MessageType::KvReq => self.store.on_request(msg.take_kv_req(), token, msg_id),
            _ => {
                Err(other(format!("unsupported message {:?} for token {:?} with msg id {}",
                                  msg_type,
                                  token,
                                  msg_id)))
            }
        }
    }

    fn on_raft_command(&mut self, msg: RaftCmdRequest, token: Token, msg_id: u64) -> Result<()> {
        debug!("handle raft command {:?}", msg);
        let ch = self.sendch.clone();
        let cb = Box::new(move |resp: RaftCmdResponse| -> RaftResult<()> {
            send_raft_cmd_resp(ch, token, msg_id, resp);
            Ok(())
        });

        let uuid = msg.get_header().get_uuid().to_vec();
        if let Err(e) = self.trans.rl().send_command(msg, cb) {
            // send error, reply an error response.
            warn!("send command for token {:?} with msg id {} err {:?}",
                  token,
                  msg_id,
                  e);

            let mut resp = cmd_resp::message_error(e);
            resp.mut_header().set_uuid(uuid);
            send_raft_cmd_resp(self.sendch.clone(), token, msg_id, resp)
        }

        Ok(())
    }

    fn on_readable(&mut self, event_loop: &mut EventLoop<Self>, token: Token) {
        match token {
            SERVER_TOKEN => {
                loop {
                    // For edge trigger, we must accept all connections until None.
                    let sock = match self.listener.accept() {
                        Err(e) => {
                            error!("accept error: {:?}", e);
                            return;
                        }
                        Ok(None) => {
                            debug!("no connection, accept later.");
                            return;
                        }
                        Ok(Some((sock, addr))) => {
                            debug!("accept conn {}", addr);
                            sock
                        }
                    };

                    if let Err(e) = self.add_new_conn(event_loop, sock, None) {
                        error!("register conn err {:?}", e);
                    }
                }
            }
            token => {
                if let Err(e) = self.on_conn_readable(event_loop, token) {
                    warn!("handle read conn for token {:?} err {:?}, remove", token, e);
                    self.remove_conn(event_loop, token);
                }
            }

        }
    }

    fn on_writable(&mut self, event_loop: &mut EventLoop<Self>, token: Token) {
        let res = match self.conns.get_mut(&token) {
            None => {
                warn!("missing conn for token {:?}", token);
                return;
            }
            Some(conn) => conn.write(event_loop),
        };

        if let Err(e) = res {
            warn!("handle write conn err {:?}, remove", e);
            self.remove_conn(event_loop, token);
        }
    }

    fn writedata(&mut self, event_loop: &mut EventLoop<Self>, token: Token, data: ConnData) {
        let res = match self.conns.get_mut(&token) {
            None => {
                warn!("missing conn for token {:?}", token);
                return;
            }
            Some(conn) => conn.append_write_buf(event_loop, data),
        };

        if let Err(e) = res {
            warn!("handle write data err {:?}, remove", e);
            self.remove_conn(event_loop, token);
        }
    }

    fn connect_peer(&mut self, event_loop: &mut EventLoop<Self>, addr: &str) -> Result<Token> {
        let peer_addr = try!(addr.parse());
        let sock = try!(TcpStream::connect(&peer_addr));
        let token = try!(self.add_new_conn(event_loop, sock, Some(addr.to_string())));
        self.peers.insert(addr.to_owned(), token);
        Ok(token)
    }

    fn sendpeer(&mut self, event_loop: &mut EventLoop<Self>, addr: String, data: ConnData) {
        // check the corresponding token for peer address.
        let mut token = self.peers.get(&addr).map_or(INVALID_TOKEN, |t| *t);

        if token == INVALID_TOKEN {
            match self.connect_peer(event_loop, &addr) {
                Err(e) => {
                    error!("connect {:?} err {:?}", addr, e);
                    return;
                }
                Ok(new_token) => token = new_token,
            }
        }

        self.writedata(event_loop, token, data);
    }
}

impl<T: Transport> Handler for Server<T> {
    type Timeout = Msg;
    type Message = Msg;

    fn ready(&mut self, event_loop: &mut EventLoop<Self>, token: Token, events: EventSet) {
        if events.is_hup() || events.is_error() {
            self.remove_conn(event_loop, token);
            return;
        }

        if events.is_readable() {
            self.on_readable(event_loop, token);
        }

        if events.is_writable() {
            self.on_writable(event_loop, token);
        }
    }

    fn notify(&mut self, event_loop: &mut EventLoop<Self>, msg: Msg) {
        match msg {
            Msg::Quit => event_loop.shutdown(),
            Msg::WriteData{token, data} => self.writedata(event_loop, token, data),
            Msg::SendPeer{addr, data} => self.sendpeer(event_loop, addr, data),
        }
    }

    fn timeout(&mut self, _: &mut EventLoop<Self>, _: Msg) {
        // nothing to do now.
    }

    fn interrupted(&mut self, event_loop: &mut EventLoop<Self>) {
        event_loop.shutdown();
    }

    fn tick(&mut self, _: &mut EventLoop<Self>) {
        // tick is called in the end of the loop, so if we notify to quit,
        // we will quit the server here.
        // TODO: handle quit server if event_loop is_running() returns false.
    }
}

fn send_raft_cmd_resp(ch: SendCh, token: Token, msg_id: u64, resp: RaftCmdResponse) {
    let mut resp_msg = Message::new();
    resp_msg.set_msg_type(MessageType::CmdResp);
    resp_msg.set_cmd_resp(resp);
    if let Err(e) = ch.send(Msg::WriteData {
        token: token,
        data: ConnData::new(msg_id, resp_msg),
    }) {
        error!("send raft cmd resp failed with token {:?}, msg id {}, err {:?}",
               token,
               msg_id,
               e);
    }
}
<commit_msg>Simplify closure<commit_after>use std::collections::HashMap;
use std::option::Option;
use std::sync::{Arc, RwLock};
use std::boxed::Box;
use std::net::SocketAddr;

use mio::{Token, Handler, EventLoop, EventSet, PollOpt};
use mio::tcp::{TcpListener, TcpStream};

use raftserver::store::{cmd_resp, Transport};
use kvproto::raft_cmdpb::{RaftCmdRequest, RaftCmdResponse};
use kvproto::msgpb::{MessageType, Message};
use super::{Msg, SendCh, ConnData};
use super::conn::Conn;
use super::{Result, other};
use util::HandyRwLock;
use storage::Storage;
use super::kv::StoreHandler;

const SERVER_TOKEN: Token = Token(1);
const FIRST_CUSTOM_TOKEN: Token = Token(1024);
const INVALID_TOKEN: Token = Token(0);

pub fn create_event_loop<T: Transport>() -> Result<EventLoop<Server<T>>> {
    let event_loop = try!(EventLoop::new());
    Ok(event_loop)
}

pub struct Server<T: Transport> {
    listener: TcpListener,
    // We use HashMap instead of common use mio slab to avoid token reusing.
    // In our raft server, a client with token 1 sends a raft command, we will
    // propose this command, execute it then send the response to the client with
    // token 1. But before the response, the client connection is broken and another
    // new client connects, mio slab may reuse the token 1 for it. So the subsequent
    // response will be sent to the new client.
    // To avoid this, we use the HashMap instead and can guarantee the token id is
    // unique and can't be reused.
    conns: HashMap<Token, Conn>,
    conn_token_counter: usize,
    sendch: SendCh,

    // addr -> Token
    peers: HashMap<String, Token>,

    trans: Arc<RwLock<T>>,

    store: StoreHandler,
}

impl<T: Transport> Server<T> {
    // Create a server with already initialized engines.
    // We must bootstrap all stores before running the server.
    pub fn new(event_loop: &mut EventLoop<Self>,
               addr: &str,
               storage: Storage,
               trans: Arc<RwLock<T>>)
               -> Result<Server<T>> {
        let laddr = try!(addr.parse());
        let listener = try!(TcpListener::bind(&laddr));

        try!(event_loop.register(&listener,
                                 SERVER_TOKEN,
                                 EventSet::readable(),
                                 PollOpt::edge()));

        let sendch = SendCh::new(event_loop.channel());

        let store_handler = StoreHandler::new(storage, sendch.clone());

        let svr = Server {
            listener: listener,
            sendch: sendch,
            conns: HashMap::new(),
            conn_token_counter: FIRST_CUSTOM_TOKEN.as_usize(),
            peers: HashMap::new(),
            trans: trans,
            store: store_handler,
        };

        Ok(svr)
    }

    pub fn get_trans(&self) -> Arc<RwLock<T>> {
        self.trans.clone()
    }

    pub fn run(&mut self, event_loop: &mut EventLoop<Self>) -> Result<()> {
        try!(event_loop.run(self));
        Ok(())
    }

    pub fn get_sendch(&self) -> SendCh {
        self.sendch.clone()
    }

    // Return listening address, this may only be used for outer test
    // to get the real address because we may use "127.0.0.1:0"
    // in test to avoid port conflict.
    pub fn listening_addr(&self) -> Result<SocketAddr> {
        let addr = try!(self.listener.local_addr());
        Ok(addr)
    }

    fn remove_conn(&mut self, event_loop: &mut EventLoop<Self>, token: Token) {
        let conn = self.conns.remove(&token);
        match conn {
            Some(conn) => {
                // if connected to remote peer, remove this too.
                if let Some(addr) = conn.peer_addr {
                    self.peers.remove(&addr);
                }

                if let Err(e) = event_loop.deregister(&conn.sock) {
                    error!("deregister conn err {:?}", e);
                }
            }
            None => {
                warn!("missing connection for token {}", token.as_usize());
            }
        }
    }

    fn add_new_conn(&mut self,
                    event_loop: &mut EventLoop<Self>,
                    sock: TcpStream,
                    peer_addr: Option<String>)
                    -> Result<Token> {
        let new_token = Token(self.conn_token_counter);
        self.conn_token_counter += 1;

        // TODO: check conn max capacity.

        try!(sock.set_nodelay(true));

        try!(event_loop.register(&sock,
                                 new_token,
                                 EventSet::readable() | EventSet::hup(),
                                 PollOpt::edge()));

        let conn = Conn::new(sock, new_token, peer_addr);
        self.conns.insert(new_token, conn);

        Ok(new_token)
    }

    fn on_conn_readable(&mut self, event_loop: &mut EventLoop<Self>, token: Token) -> Result<()> {
        let msgs = try!(match self.conns.get_mut(&token) {
            None => {
                warn!("missing conn for token {:?}", token);
                return Ok(());
            }
            Some(conn) => conn.read(event_loop),
        });

        if msgs.is_empty() {
            // Read no message, no need to handle.
            return Ok(());
        }

        for msg in msgs {
            try!(self.on_conn_msg(token, msg))
        }

        Ok(())
    }

    fn on_conn_msg(&mut self, token: Token, data: ConnData) -> Result<()> {
        let msg_id = data.msg_id;
        let mut msg = data.msg;

        let msg_type = msg.get_msg_type();
        match msg_type {
            MessageType::Raft => {
                if let Err(e) = self.trans.rl().send_raft_msg(msg.take_raft()) {
                    // Should we return error to let outer close this connection later?
                    error!("send raft message for token {:?} with msg id {} err {:?}",
                           token,
                           msg_id,
                           e);
                }
                Ok(())
            }
            MessageType::Cmd => self.on_raft_command(msg.take_cmd_req(), token, msg_id),
            MessageType::KvReq => self.store.on_request(msg.take_kv_req(), token, msg_id),
            _ => {
                Err(other(format!("unsupported message {:?} for token {:?} with msg id {}",
                                  msg_type,
                                  token,
                                  msg_id)))
            }
        }
    }

    fn on_raft_command(&mut self, msg: RaftCmdRequest, token: Token, msg_id: u64) -> Result<()> {
        debug!("handle raft command {:?}", msg);
        let ch = self.sendch.clone();
        let cb = Box::new(move |resp| {
            send_raft_cmd_resp(ch, token, msg_id, resp);
            Ok(())
        });

        let uuid = msg.get_header().get_uuid().to_vec();
        if let Err(e) = self.trans.rl().send_command(msg, cb) {
            // send error, reply an error response.
            warn!("send command for token {:?} with msg id {} err {:?}",
                  token,
                  msg_id,
                  e);

            let mut resp = cmd_resp::message_error(e);
            resp.mut_header().set_uuid(uuid);
            send_raft_cmd_resp(self.sendch.clone(), token, msg_id, resp)
        }

        Ok(())
    }

    fn on_readable(&mut self, event_loop: &mut EventLoop<Self>, token: Token) {
        match token {
            SERVER_TOKEN => {
                loop {
                    // For edge trigger, we must accept all connections until None.
                    let sock = match self.listener.accept() {
                        Err(e) => {
                            error!("accept error: {:?}", e);
                            return;
                        }
                        Ok(None) => {
                            debug!("no connection, accept later.");
                            return;
                        }
                        Ok(Some((sock, addr))) => {
                            debug!("accept conn {}", addr);
                            sock
                        }
                    };

                    if let Err(e) = self.add_new_conn(event_loop, sock, None) {
                        error!("register conn err {:?}", e);
                    }
                }
            }
            token => {
                if let Err(e) = self.on_conn_readable(event_loop, token) {
                    warn!("handle read conn for token {:?} err {:?}, remove", token, e);
                    self.remove_conn(event_loop, token);
                }
            }

        }
    }

    fn on_writable(&mut self, event_loop: &mut EventLoop<Self>, token: Token) {
        let res = match self.conns.get_mut(&token) {
            None => {
                warn!("missing conn for token {:?}", token);
                return;
            }
            Some(conn) => conn.write(event_loop),
        };

        if let Err(e) = res {
            warn!("handle write conn err {:?}, remove", e);
            self.remove_conn(event_loop, token);
        }
    }

    fn writedata(&mut self, event_loop: &mut EventLoop<Self>, token: Token, data: ConnData) {
        let res = match self.conns.get_mut(&token) {
            None => {
                warn!("missing conn for token {:?}", token);
                return;
            }
            Some(conn) => conn.append_write_buf(event_loop, data),
        };

        if let Err(e) = res {
            warn!("handle write data err {:?}, remove", e);
            self.remove_conn(event_loop, token);
        }
    }

    fn connect_peer(&mut self, event_loop: &mut EventLoop<Self>, addr: &str) -> Result<Token> {
        let peer_addr = try!(addr.parse());
        let sock = try!(TcpStream::connect(&peer_addr));
        let token = try!(self.add_new_conn(event_loop, sock, Some(addr.to_string())));
        self.peers.insert(addr.to_owned(), token);
        Ok(token)
    }

    fn sendpeer(&mut self, event_loop: &mut EventLoop<Self>, addr: String, data: ConnData) {
        // check the corresponding token for peer address.
        let mut token = self.peers.get(&addr).map_or(INVALID_TOKEN, |t| *t);

        if token == INVALID_TOKEN {
            match self.connect_peer(event_loop, &addr) {
                Err(e) => {
                    error!("connect {:?} err {:?}", addr, e);
                    return;
                }
                Ok(new_token) => token = new_token,
            }
        }

        self.writedata(event_loop, token, data);
    }
}

impl<T: Transport> Handler for Server<T> {
    type Timeout = Msg;
    type Message = Msg;

    fn ready(&mut self, event_loop: &mut EventLoop<Self>, token: Token, events: EventSet) {
        if events.is_hup() || events.is_error() {
            self.remove_conn(event_loop, token);
            return;
        }

        if events.is_readable() {
            self.on_readable(event_loop, token);
        }

        if events.is_writable() {
            self.on_writable(event_loop, token);
        }
    }

    fn notify(&mut self, event_loop: &mut EventLoop<Self>, msg: Msg) {
        match msg {
            Msg::Quit => event_loop.shutdown(),
            Msg::WriteData{token, data} => self.writedata(event_loop, token, data),
            Msg::SendPeer{addr, data} => self.sendpeer(event_loop, addr, data),
        }
    }

    fn timeout(&mut self, _: &mut EventLoop<Self>, _: Msg) {
        // nothing to do now.
    }

    fn interrupted(&mut self, event_loop: &mut EventLoop<Self>) {
        event_loop.shutdown();
    }

    fn tick(&mut self, _: &mut EventLoop<Self>) {
        // tick is called in the end of the loop, so if we notify to quit,
        // we will quit the server here.
        // TODO: handle quit server if event_loop is_running() returns false.
    }
}

fn send_raft_cmd_resp(ch: SendCh, token: Token, msg_id: u64, resp: RaftCmdResponse) {
    let mut resp_msg = Message::new();
    resp_msg.set_msg_type(MessageType::CmdResp);
    resp_msg.set_cmd_resp(resp);
    if let Err(e) = ch.send(Msg::WriteData {
        token: token,
        data: ConnData::new(msg_id, resp_msg),
    }) {
        error!("send raft cmd resp failed with token {:?}, msg id {}, err {:?}",
               token,
               msg_id,
               e);
    }
}
<|endoftext|>


<commit_before>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io

/// In this module we put the code related to stops, that need to draw on 'places', 'mimir',
/// 'common', and 'config' (ie all the workspaces that make up mimirsbrunn).
use futures::stream::{Stream, TryStreamExt};
use mimir::domain::model::configuration::{ContainerConfig, PhysicalModeWeight};
use snafu::{ResultExt, Snafu};
use std::collections::HashMap;
use std::ops::Deref;
use std::path::PathBuf;
use std::sync::Arc;
use tracing::{info, warn};

use crate::admin_geofinder::AdminGeoFinder;
use crate::labels;
use mimir::adapters::secondary::elasticsearch::{self, ElasticsearchStorage};
use mimir::domain::ports::primary::{generate_index::GenerateIndex, list_documents::ListDocuments};
use places::admin::Admin;
use places::stop::Stop;

#[derive(Debug, Snafu)]
pub enum Error {
    // #[snafu(display("Settings (Configuration or CLI) Error: {}", source))]
    // Settings { source: settings::Error },
    #[snafu(display("Elasticsearch Connection Pool {}", source))]
    ElasticsearchPool {
        source: elasticsearch::remote::Error,
    },

    #[snafu(display("Elasticsearch Connection Pool {}", source))]
    ElasticsearchConnection {
        source: mimir::domain::ports::secondary::remote::Error,
    },

    // #[snafu(display("Cosmogony Error: {}", details))]
    // Cosmogony { details: String },
    #[snafu(display("Index Generation Error {}", source))]
    IndexGeneration {
        source: mimir::domain::model::error::Error,
    },

    // transit_model uses failure::Error, which does not implement std::Error, so
    // we use a String to get the error message instead.
    #[snafu(display("Transit Model Error {}", details))]
    TransitModel { details: String },

    #[snafu(display("Admin Retrieval Error {}", details))]
    AdminRetrieval { details: String },
}

pub fn initialize_weights<'a, It, S: ::std::hash::BuildHasher>(
    stops: It,
    nb_stop_points: &HashMap<String, u32, S>,
) where
    It: Iterator<Item = &'a mut Stop>,
{
    let max = *nb_stop_points.values().max().unwrap_or(&1) as f64;
    for stop in stops {
        stop.weight = if let Some(weight) = nb_stop_points.get(&stop.id) {
            *weight as f64 / max
        } else {
            0.0
        };
    }
}

pub fn build_stop_area_weight(
    navitia: &transit_model::Model,
    physical_mode_weight: &Option<Vec<PhysicalModeWeight>>,
) -> HashMap<String, f64> {
    let md_weight_hash_map: HashMap<String, f64> = match physical_mode_weight {
        Some(modes) => modes
            .iter()
            .map(|mode| (mode.id.to_string().to_lowercase(), mode.weight as f64))
            .collect::<HashMap<String, f64>>(),
        _ => HashMap::new(),
    };
    let mut result = HashMap::new();
    for (idx, sp) in &navitia.stop_points {
        let stop_area = places::utils::normalize_id("stop_area", &sp.stop_area_id);

        let mut weights: f64 = navitia
            .get_corresponding_from_idx::<_, transit_model::objects::PhysicalMode>(idx)
            .into_iter()
            .map(|idx| {
                let ph_mode = &navitia.physical_modes[idx].id;
                let pm_w = md_weight_hash_map.get(&ph_mode.to_lowercase());
                match pm_w {
                    Some(value) => *value,
                    _ => {
                        warn!("Physical mode, id: {}, not found in mimir config.", ph_mode);
                        0.0
                    }
                }
            })
            .filter(|&weight| weight != 0.0)
            .into_iter()
            .sum();
        let res = result.get(&stop_area);
        if let Some(value) = res {
            weights += value;
        }
        result.insert(stop_area, weights);
    }
    result
}

pub fn make_weight(stop: &mut Stop, stop_areas_weights: &HashMap<String, f64>) {
    // Admin weight
    let mut admin_weight = stop
        .administrative_regions
        .iter()
        .filter(|adm| adm.is_city())
        .map(|adm| adm.weight)
        .next()
        .unwrap_or(0.0);
    // FIXME: 1024, automagic!
    // It's a factor used to bring the stop weight and the admin weight in the same order of
    // magnitude...
    // We then use a log to compress the distance between low admin weight and high ones.
    admin_weight = admin_weight * 1024.0 + 1.0;
    admin_weight = admin_weight.log10();

    let weights = stop_areas_weights.get(&*stop.id);
    if let Some(value) = weights {
        stop.weight = (value + admin_weight) / (2 as f64);
    } else {
        stop.weight = admin_weight
    }
}

fn attach_stop(stop: &mut Stop, admins: Vec<Arc<Admin>>) {
    let admins_iter = admins.iter().map(|a| a.deref());
    let country_codes = places::admin::find_country_codes(admins_iter.clone());

    stop.label = labels::format_stop_label(&stop.name, admins_iter, &country_codes);
    stop.zip_codes = places::admin::get_zip_codes_from_admins(&admins);

    stop.country_codes = country_codes;
    stop.administrative_regions = admins;
}

/// Attach the stops to administrative regions
///
/// The admins are loaded from Elasticsearch and stored in a quadtree
/// We attach a stop with all the admins that have a boundary containing
/// the coordinate of the stop
async fn attach_stops_to_admins<'a, It: Iterator<Item = &'a mut Stop>>(
    stops: It,
    client: &ElasticsearchStorage,
) -> Result<(), Error> {
    match client.list_documents().await {
        Ok(stream) => {
            let admins: Vec<Admin> = stream.try_collect().await.context(IndexGenerationSnafu)?;

            if admins.is_empty() {
                return Err(Error::AdminRetrieval {
                    details: String::from("no admin retrieved to enrich stops"),
                });
            }
            info!("{} admins retrieved from ES ", admins.len());
            let admins_geofinder = admins.into_iter().collect::<AdminGeoFinder>();

            let mut nb_unmatched = 0u32;
            let mut nb_matched = 0u32;
            // FIXME Opportunity for concurrent work
            for mut stop in stops {
                let admins = admins_geofinder.get(&stop.coord);

                if admins.is_empty() {
                    nb_unmatched += 1;
                } else {
                    nb_matched += 1;
                }

                attach_stop(&mut stop, admins);
            }

            info!(
                "there are {}/{} stops without any admin",
                nb_unmatched,
                nb_matched + nb_unmatched
            );
            Ok(())
        }
        Err(_) => Err(Error::AdminRetrieval {
            details: String::from("Could not retrieve admins to enrich stops"),
        }),
    }
}

/// Stores the stops found in the 'input' directory, in Elasticsearch, with the given
/// configuration.
///
/// The main part of this function is to actually create a list of stops
/// from the information found in the NTFS directory.
pub async fn index_ntfs(
    input: PathBuf,
    config: &ContainerConfig,
    physical_mode_weight: &Option<Vec<PhysicalModeWeight>>,
    client: &ElasticsearchStorage,
) -> Result<(), Error> {
    let navitia = transit_model::ntfs::read(&input).map_err(|err| Error::TransitModel {
        details: format!(
            "Could not read transit model from {}: {}",
            input.display(),
            err.to_string()
        ),
    })?;

    info!("Build stops weight by physical modes");
    let stop_areas_weights = build_stop_area_weight(&navitia, &physical_mode_weight);

    info!("Make mimir stops from navitia stops");
    let mut stops: Vec<Stop> = navitia
        .stop_areas
        .iter()
        .map(|(idx, sa)| places::stop::to_mimir(idx, sa, &navitia))
        .collect();

    info!("Attach stops to admins");
    attach_stops_to_admins(stops.iter_mut(), client).await?;

    for stop in &mut stops {
        stop.coverages.push(config.dataset.clone());
        make_weight(stop, &stop_areas_weights);
    }

    tracing::info!("Beginning to import stops into elasticsearch.");
    import_stops(client, config, futures::stream::iter(stops)).await
}

// FIXME Should not be ElasticsearchStorage, but rather a trait GenerateIndex
pub async fn import_stops<S>(
    client: &ElasticsearchStorage,
    config: &ContainerConfig,
    stops: S,
) -> Result<(), Error>
where
    S: Stream<Item = Stop> + Send + Sync + Unpin + 'static,
{
    client
        .generate_index(config, stops)
        .await
        .context(IndexGenerationSnafu)?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use crate::stops::make_weight;
    use cosmogony::ZoneType;
    use places::admin::Admin;
    use places::stop::Stop;
    use serial_test::serial;
    use std::collections::HashMap;
    use std::sync::Arc;

    fn approx_equal(a: f64, b: f64, dp: u8) -> bool {
        let p = 10f64.powi(-(dp as i32));
        (a - b).abs() < p
    }

    #[tokio::test]
    #[serial]
    async fn test_make_weight_without_physical_mode_weight_and_without_admins() {
        let mut stop = Stop {
            id: "123".to_string(),
            ..Default::default()
        };
        make_weight(&mut stop, &HashMap::new());
        assert!(approx_equal(stop.weight, 0.0, 1));
    }

    #[tokio::test]
    #[serial]
    async fn test_make_weight_without_physical_mode_weight_and_with_admins() {
        let admin = Admin {
            id: "adm:01".to_string(),
            weight: 0.12,
            zone_type: Some(ZoneType::City),
            ..Default::default()
        };
        let mut stop = Stop {
            id: "123".to_string(),
            administrative_regions: vec![Arc::new(admin)],
            ..Default::default()
        };
        make_weight(&mut stop, &HashMap::new());
        approx_equal(stop.weight, 2.0930, 4);
    }

    #[tokio::test]
    #[serial]
    async fn test_make_weight_with_physical_mode_weight_and_with_admins() {
        let mut physical_mode_weight = HashMap::new();
        physical_mode_weight.insert("123".to_string(), 5.0);

        let admin = Admin {
            id: "adm:01".to_string(),
            weight: 0.12,
            zone_type: Some(ZoneType::City),
            ..Default::default()
        };
        let mut stop = Stop {
            id: "123".to_string(),
            administrative_regions: vec![Arc::new(admin)],
            ..Default::default()
        };
        make_weight(&mut stop, &physical_mode_weight);
        approx_equal(stop.weight, 3.5465, 4);
    }
}
<commit_msg>Remove log10 in weight stops object<commit_after>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io

/// In this module we put the code related to stops, that need to draw on 'places', 'mimir',
/// 'common', and 'config' (ie all the workspaces that make up mimirsbrunn).
use futures::stream::{Stream, TryStreamExt};
use mimir::domain::model::configuration::{ContainerConfig, PhysicalModeWeight};
use snafu::{ResultExt, Snafu};
use std::collections::HashMap;
use std::ops::Deref;
use std::path::PathBuf;
use std::sync::Arc;
use tracing::{info, warn};

use crate::admin_geofinder::AdminGeoFinder;
use crate::labels;
use mimir::adapters::secondary::elasticsearch::{self, ElasticsearchStorage};
use mimir::domain::ports::primary::{generate_index::GenerateIndex, list_documents::ListDocuments};
use places::admin::Admin;
use places::stop::Stop;

#[derive(Debug, Snafu)]
pub enum Error {
    // #[snafu(display("Settings (Configuration or CLI) Error: {}", source))]
    // Settings { source: settings::Error },
    #[snafu(display("Elasticsearch Connection Pool {}", source))]
    ElasticsearchPool {
        source: elasticsearch::remote::Error,
    },

    #[snafu(display("Elasticsearch Connection Pool {}", source))]
    ElasticsearchConnection {
        source: mimir::domain::ports::secondary::remote::Error,
    },

    // #[snafu(display("Cosmogony Error: {}", details))]
    // Cosmogony { details: String },
    #[snafu(display("Index Generation Error {}", source))]
    IndexGeneration {
        source: mimir::domain::model::error::Error,
    },

    // transit_model uses failure::Error, which does not implement std::Error, so
    // we use a String to get the error message instead.
    #[snafu(display("Transit Model Error {}", details))]
    TransitModel { details: String },

    #[snafu(display("Admin Retrieval Error {}", details))]
    AdminRetrieval { details: String },
}

pub fn initialize_weights<'a, It, S: ::std::hash::BuildHasher>(
    stops: It,
    nb_stop_points: &HashMap<String, u32, S>,
) where
    It: Iterator<Item = &'a mut Stop>,
{
    let max = *nb_stop_points.values().max().unwrap_or(&1) as f64;
    for stop in stops {
        stop.weight = if let Some(weight) = nb_stop_points.get(&stop.id) {
            *weight as f64 / max
        } else {
            0.0
        };
    }
}

pub fn build_stop_area_weight(
    navitia: &transit_model::Model,
    physical_mode_weight: &Option<Vec<PhysicalModeWeight>>,
) -> HashMap<String, f64> {
    let md_weight_hash_map: HashMap<String, f64> = match physical_mode_weight {
        Some(modes) => modes
            .iter()
            .map(|mode| (mode.id.to_string().to_lowercase(), mode.weight as f64))
            .collect::<HashMap<String, f64>>(),
        _ => HashMap::new(),
    };
    let mut result = HashMap::new();
    for (idx, sp) in &navitia.stop_points {
        let stop_area = places::utils::normalize_id("stop_area", &sp.stop_area_id);

        let mut weights: f64 = navitia
            .get_corresponding_from_idx::<_, transit_model::objects::PhysicalMode>(idx)
            .into_iter()
            .map(|idx| {
                let ph_mode = &navitia.physical_modes[idx].id;
                let pm_w = md_weight_hash_map.get(&ph_mode.to_lowercase());
                match pm_w {
                    Some(value) => *value,
                    _ => {
                        warn!("Physical mode, id: {}, not found in mimir config.", ph_mode);
                        0.0
                    }
                }
            })
            .filter(|&weight| weight != 0.0)
            .into_iter()
            .sum();
        let res = result.get(&stop_area);
        if let Some(value) = res {
            weights += value;
        }
        result.insert(stop_area, weights);
    }
    result
}

pub fn make_weight(stop: &mut Stop, stop_areas_weights: &HashMap<String, f64>) {
    // Admin weight
    let admin_weight = stop
        .administrative_regions
        .iter()
        .filter(|adm| adm.is_city())
        .map(|adm| adm.weight)
        .next()
        .unwrap_or(0.0);
    // FIXME: 1024, automagic!
    // It's a factor used to bring the stop weight and the admin weight in the same order of
    // magnitude...
    // We then use a log to compress the distance between low admin weight and high ones.
    // admin_weight = admin_weight * 1024.0 + 1.0;
    // admin_weight = admin_weight.log10();

    let weights = stop_areas_weights.get(&*stop.id);
    if let Some(value) = weights {
        stop.weight = (value + admin_weight) / (2 as f64);
    } else {
        stop.weight = admin_weight
    }
}

fn attach_stop(stop: &mut Stop, admins: Vec<Arc<Admin>>) {
    let admins_iter = admins.iter().map(|a| a.deref());
    let country_codes = places::admin::find_country_codes(admins_iter.clone());

    stop.label = labels::format_stop_label(&stop.name, admins_iter, &country_codes);
    stop.zip_codes = places::admin::get_zip_codes_from_admins(&admins);

    stop.country_codes = country_codes;
    stop.administrative_regions = admins;
}

/// Attach the stops to administrative regions
///
/// The admins are loaded from Elasticsearch and stored in a quadtree
/// We attach a stop with all the admins that have a boundary containing
/// the coordinate of the stop
async fn attach_stops_to_admins<'a, It: Iterator<Item = &'a mut Stop>>(
    stops: It,
    client: &ElasticsearchStorage,
) -> Result<(), Error> {
    match client.list_documents().await {
        Ok(stream) => {
            let admins: Vec<Admin> = stream.try_collect().await.context(IndexGenerationSnafu)?;

            if admins.is_empty() {
                return Err(Error::AdminRetrieval {
                    details: String::from("no admin retrieved to enrich stops"),
                });
            }
            info!("{} admins retrieved from ES ", admins.len());
            let admins_geofinder = admins.into_iter().collect::<AdminGeoFinder>();

            let mut nb_unmatched = 0u32;
            let mut nb_matched = 0u32;
            // FIXME Opportunity for concurrent work
            for mut stop in stops {
                let admins = admins_geofinder.get(&stop.coord);

                if admins.is_empty() {
                    nb_unmatched += 1;
                } else {
                    nb_matched += 1;
                }

                attach_stop(&mut stop, admins);
            }

            info!(
                "there are {}/{} stops without any admin",
                nb_unmatched,
                nb_matched + nb_unmatched
            );
            Ok(())
        }
        Err(_) => Err(Error::AdminRetrieval {
            details: String::from("Could not retrieve admins to enrich stops"),
        }),
    }
}

/// Stores the stops found in the 'input' directory, in Elasticsearch, with the given
/// configuration.
///
/// The main part of this function is to actually create a list of stops
/// from the information found in the NTFS directory.
pub async fn index_ntfs(
    input: PathBuf,
    config: &ContainerConfig,
    physical_mode_weight: &Option<Vec<PhysicalModeWeight>>,
    client: &ElasticsearchStorage,
) -> Result<(), Error> {
    let navitia = transit_model::ntfs::read(&input).map_err(|err| Error::TransitModel {
        details: format!(
            "Could not read transit model from {}: {}",
            input.display(),
            err.to_string()
        ),
    })?;

    info!("Build stops weight by physical modes");
    let stop_areas_weights = build_stop_area_weight(&navitia, &physical_mode_weight);

    info!("Make mimir stops from navitia stops");
    let mut stops: Vec<Stop> = navitia
        .stop_areas
        .iter()
        .map(|(idx, sa)| places::stop::to_mimir(idx, sa, &navitia))
        .collect();

    info!("Attach stops to admins");
    attach_stops_to_admins(stops.iter_mut(), client).await?;

    for stop in &mut stops {
        stop.coverages.push(config.dataset.clone());
        make_weight(stop, &stop_areas_weights);
    }

    tracing::info!("Beginning to import stops into elasticsearch.");
    import_stops(client, config, futures::stream::iter(stops)).await
}

// FIXME Should not be ElasticsearchStorage, but rather a trait GenerateIndex
pub async fn import_stops<S>(
    client: &ElasticsearchStorage,
    config: &ContainerConfig,
    stops: S,
) -> Result<(), Error>
where
    S: Stream<Item = Stop> + Send + Sync + Unpin + 'static,
{
    client
        .generate_index(config, stops)
        .await
        .context(IndexGenerationSnafu)?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use crate::stops::make_weight;
    use cosmogony::ZoneType;
    use places::admin::Admin;
    use places::stop::Stop;
    use serial_test::serial;
    use std::collections::HashMap;
    use std::sync::Arc;

    fn approx_equal(a: f64, b: f64, dp: u8) -> bool {
        let p = 10f64.powi(-(dp as i32));
        (a - b).abs() < p
    }

    #[tokio::test]
    #[serial]
    async fn test_make_weight_without_physical_mode_weight_and_without_admins() {
        let mut stop = Stop {
            id: "123".to_string(),
            ..Default::default()
        };
        make_weight(&mut stop, &HashMap::new());
        assert!(approx_equal(stop.weight, 0.0, 1));
    }

    #[tokio::test]
    #[serial]
    async fn test_make_weight_without_physical_mode_weight_and_with_admins() {
        let admin = Admin {
            id: "adm:01".to_string(),
            weight: 0.12,
            zone_type: Some(ZoneType::City),
            ..Default::default()
        };
        let mut stop = Stop {
            id: "123".to_string(),
            administrative_regions: vec![Arc::new(admin)],
            ..Default::default()
        };
        make_weight(&mut stop, &HashMap::new());
        approx_equal(stop.weight, 2.0930, 4);
    }

    #[tokio::test]
    #[serial]
    async fn test_make_weight_with_physical_mode_weight_and_with_admins() {
        let mut physical_mode_weight = HashMap::new();
        physical_mode_weight.insert("123".to_string(), 5.0);

        let admin = Admin {
            id: "adm:01".to_string(),
            weight: 0.12,
            zone_type: Some(ZoneType::City),
            ..Default::default()
        };
        let mut stop = Stop {
            id: "123".to_string(),
            administrative_regions: vec![Arc::new(admin)],
            ..Default::default()
        };
        make_weight(&mut stop, &physical_mode_weight);
        approx_equal(stop.weight, 3.5465, 4);
    }
}
<|endoftext|>


<commit_before>include!("header.rs");

macro_rules! impl_test {
    ($modname:ident, $random:path) => {
mod $modname {
    use ndarray::prelude::*;
    use ndarray_linalg::prelude::*;
    use ndarray_numtest::prelude::*;
    use ndarray_rand::RandomExt;
    #[test]
    fn solve_upper() {
        let r_dist = RealNormal::new(0.0, 1.0);
        let a = drop_lower($random((3, 3), r_dist));
        println!("a = \n{:?}", &a);
        let b = $random(3, r_dist);
        println!("b = \n{:?}", &b);
        let x = a.solve_upper(b.clone()).unwrap();
        println!("x = \n{:?}", &x);
        println!("Ax = \n{:?}", a.dot(&x));
        all_close_l2(&a.dot(&x), &b, 1e-7).unwrap();
    }

    #[test]
    fn solve_upper_t() {
        let r_dist = RealNormal::new(0., 1.);
        let a = drop_lower($random((3, 3), r_dist).reversed_axes());
        println!("a = \n{:?}", &a);
        let b = $random(3, r_dist);
        println!("b = \n{:?}", &b);
        let x = a.solve_upper(b.clone()).unwrap();
        println!("x = \n{:?}", &x);
        println!("Ax = \n{:?}", a.dot(&x));
        all_close_l2(&a.dot(&x), &b, 1e-7).unwrap();
    }

    #[test]
    fn solve_lower() {
        let r_dist = RealNormal::new(0., 1.);
        let a = drop_upper($random((3, 3), r_dist));
        println!("a = \n{:?}", &a);
        let b = $random(3, r_dist);
        println!("b = \n{:?}", &b);
        let x = a.solve_lower(b.clone()).unwrap();
        println!("x = \n{:?}", &x);
        println!("Ax = \n{:?}", a.dot(&x));
        all_close_l2(&a.dot(&x), &b, 1e-7).unwrap();
    }

    #[test]
    fn solve_lower_t() {
        let r_dist = RealNormal::new(0., 1.);
        let a = drop_upper($random((3, 3), r_dist).reversed_axes());
        println!("a = \n{:?}", &a);
        let b = $random(3, r_dist);
        println!("b = \n{:?}", &b);
        let x = a.solve_lower(b.clone()).unwrap();
        println!("x = \n{:?}", &x);
        println!("Ax = \n{:?}", a.dot(&x));
        all_close_l2(&a.dot(&x), &b, 1e-7).unwrap();
    }
}
}} // impl_test_opnorm

impl_test!(owned, Array<f64, _>::random);
impl_test!(shared, RcArray<f64, _>::random);
<commit_msg>Add test for 2d SolveTriangular<commit_after>include!("header.rs");

macro_rules! impl_test {
    ($modname:ident, $random:path) => {
mod $modname {
    use ndarray::prelude::*;
    use ndarray_linalg::prelude::*;
    use ndarray_numtest::prelude::*;
    use ndarray_rand::RandomExt;
    #[test]
    fn solve_upper() {
        let r_dist = RealNormal::new(0.0, 1.0);
        let a = drop_lower($random((3, 3), r_dist));
        println!("a = \n{:?}", &a);
        let b = $random(3, r_dist);
        println!("b = \n{:?}", &b);
        let x = a.solve_upper(b.clone()).unwrap();
        println!("x = \n{:?}", &x);
        println!("Ax = \n{:?}", a.dot(&x));
        all_close_l2(&a.dot(&x), &b, 1e-7).unwrap();
    }

    #[test]
    fn solve_upper_t() {
        let r_dist = RealNormal::new(0., 1.);
        let a = drop_lower($random((3, 3), r_dist).reversed_axes());
        println!("a = \n{:?}", &a);
        let b = $random(3, r_dist);
        println!("b = \n{:?}", &b);
        let x = a.solve_upper(b.clone()).unwrap();
        println!("x = \n{:?}", &x);
        println!("Ax = \n{:?}", a.dot(&x));
        all_close_l2(&a.dot(&x), &b, 1e-7).unwrap();
    }

    #[test]
    fn solve_lower() {
        let r_dist = RealNormal::new(0., 1.);
        let a = drop_upper($random((3, 3), r_dist));
        println!("a = \n{:?}", &a);
        let b = $random(3, r_dist);
        println!("b = \n{:?}", &b);
        let x = a.solve_lower(b.clone()).unwrap();
        println!("x = \n{:?}", &x);
        println!("Ax = \n{:?}", a.dot(&x));
        all_close_l2(&a.dot(&x), &b, 1e-7).unwrap();
    }

    #[test]
    fn solve_lower_t() {
        let r_dist = RealNormal::new(0., 1.);
        let a = drop_upper($random((3, 3), r_dist).reversed_axes());
        println!("a = \n{:?}", &a);
        let b = $random(3, r_dist);
        println!("b = \n{:?}", &b);
        let x = a.solve_lower(b.clone()).unwrap();
        println!("x = \n{:?}", &x);
        println!("Ax = \n{:?}", a.dot(&x));
        all_close_l2(&a.dot(&x), &b, 1e-7).unwrap();
    }
}
}} // impl_test_opnorm

impl_test!(owned, Array<f64, _>::random);
impl_test!(shared, RcArray<f64, _>::random);

macro_rules! impl_test_2d {
    ($modname:ident, $drop:path, $solve:ident) => {
mod $modname {
    use super::random_owned;
    use ndarray_linalg::prelude::*;
    #[test]
    fn solve_tt() {
        let a = $drop(random_owned(3, 3, true));
        println!("a = \n{:?}", &a);
        let b = random_owned(3, 2, true);
        println!("b = \n{:?}", &b);
        let x = a.$solve(&b).unwrap();
        println!("x = \n{:?}", &x);
        println!("Ax = \n{:?}", a.dot(&x));
        all_close_l2(&a.dot(&x), &b, 1e-7).unwrap();
    }
    #[test]
    fn solve_tf() {
        let a = $drop(random_owned(3, 3, true));
        println!("a = \n{:?}", &a);
        let b = random_owned(3, 2, false);
        println!("b = \n{:?}", &b);
        let x = a.$solve(&b).unwrap();
        println!("x = \n{:?}", &x);
        println!("Ax = \n{:?}", a.dot(&x));
        all_close_l2(&a.dot(&x), &b, 1e-7).unwrap();
    }
    #[test]
    fn solve_ft() {
        let a = $drop(random_owned(3, 3, false));
        println!("a = \n{:?}", &a);
        let b = random_owned(3, 2, true);
        println!("b = \n{:?}", &b);
        let x = a.$solve(&b).unwrap();
        println!("x = \n{:?}", &x);
        println!("Ax = \n{:?}", a.dot(&x));
        all_close_l2(&a.dot(&x), &b, 1e-7).unwrap();
    }
    #[test]
    fn solve_ff() {
        let a = $drop(random_owned(3, 3, false));
        println!("a = \n{:?}", &a);
        let b = random_owned(3, 2, false);
        println!("b = \n{:?}", &b);
        let x = a.$solve(&b).unwrap();
        println!("x = \n{:?}", &x);
        println!("Ax = \n{:?}", a.dot(&x));
        all_close_l2(&a.dot(&x), &b, 1e-7).unwrap();
    }
}
}} // impl_test_2d

impl_test_2d!(lower2d, drop_upper, solve_lower);
impl_test_2d!(upper2d, drop_lower, solve_upper);
<|endoftext|>


<commit_before>//! Some unit tests for all the internal functionality of lrs

#[cfg(test)]
#[allow(dead_code)]
#[allow(unused_macros)]
#[allow(unused_imports)]
mod tests {
    use terms::*;
    use clause::*;

    #[test]
     fn lrs_symbol_create() {
        let s = symbol!["A"];
        assert!(s.val == 'A');
        assert!(s.state == true);
    }

    #[test]
    fn lrs_term_create() {
        let term = term!["A", "!🐎", "🐟", "!🚀"];

        assert!(term.contains(&symbol!["A"]));
        assert!(term.contains(&symbol!["!🐎"]));
        assert!(term.contains(&symbol!["🐟"]));
        assert!(term.contains(&symbol!["!🚀"]));
    }

    #[test]
    fn lrs_term_remove() {
        let mut term = term!["A", "!🐎", "🐟", "!🚀"];

        term.remove(symbol!["A"]);
        assert!(! term.contains(&symbol!["A"]));

        term.remove(symbol!["!🐎"]);
        assert!(! term.contains(&symbol!["!🐎"]));

        term.remove(symbol!["🐟"]);
        assert!(! term.contains(&symbol!["🐟"]));

        term.remove(symbol!["!🚀"]);
        assert!(! term.contains(&symbol!["!🚀"]));
    }

    #[test]
    fn lrs_compare_terms() {
        let a = term!["A", "B"];
        let b = term!["B", "C"];
        let clause = clause![a, b];

        assert!(clause.contains(&term!["A", "B"]));
        assert!(clause.contains(&term!["B", "A"]));

        assert!(clause.contains(&term!["C", "B"]));
        assert!(! clause.contains(&term!["C", "D"]));
    }

    #[test]
    fn lrs_simple_reduce() {
        let a = term!["A", "!B"];
        let b = term!["C", "!D"];
        let c = term!["E"];

        let mut clause = clause![a, b, c];
        clause.reduce();

        /* We should no longer contain "E" */
        assert!(! clause.contains(&term!["E"]));
    }
}
<commit_msg>Adding unit test that compares clauses<commit_after>//! Some unit tests for all the internal functionality of lrs

#[cfg(test)]
#[allow(dead_code)]
#[allow(unused_macros)]
#[allow(unused_imports)]
mod tests {
    use terms::*;
    use clause::*;

    #[test]
     fn lrs_symbol_create() {
        let s = symbol!["A"];
        assert!(s.val == 'A');
        assert!(s.state == true);
    }

    #[test]
    fn lrs_term_create() {
        let term = term!["A", "!🐎", "🐟", "!🚀"];

        assert!(term.contains(&symbol!["A"]));
        assert!(term.contains(&symbol!["!🐎"]));
        assert!(term.contains(&symbol!["🐟"]));
        assert!(term.contains(&symbol!["!🚀"]));
    }

    #[test]
    fn lrs_term_remove() {
        let mut term = term!["A", "!🐎", "🐟", "!🚀"];

        term.remove(symbol!["A"]);
        assert!(! term.contains(&symbol!["A"]));

        term.remove(symbol!["!🐎"]);
        assert!(! term.contains(&symbol!["!🐎"]));

        term.remove(symbol!["🐟"]);
        assert!(! term.contains(&symbol!["🐟"]));

        term.remove(symbol!["!🚀"]);
        assert!(! term.contains(&symbol!["!🚀"]));
    }

    #[test]
    fn lrs_compare_terms() {
        let a = term!["A", "B"];
        let b = term!["B", "C"];
        let clause = clause![a, b];

        assert!(clause.contains(&term!["A", "B"]));
        assert!(clause.contains(&term!["B", "A"]));

        assert!(clause.contains(&term!["C", "B"]));
        assert!(! clause.contains(&term!["C", "D"]));
    }

    #[test]
    fn lrs_compare_clauses() {
        let c1 = clause![term!["A"], term!["!B"]];
        let c2 = clause![term!["A"], term!["!B"]];
        assert_eq!(c1, c2);

        let c3 = clause![term!["A", "B"], term!["!B"]];
        let c4 = clause![term!["B", "A"], term!["!B"]];
        assert_eq!(c3, c4);

        let c5 = clause![term!["A", "B"], term!["B"]];
        let c6 = clause![term!["A"], term!["!B"]];
        assert!(c5 != c6);
    }

    #[test]
    fn lrs_simple_reduce() {
        let a = term!["A", "!B"];
        let b = term!["C", "!D"];
        let c = term!["E"];

        let mut clause = clause![a, b, c];
        clause.reduce();

        /* We should no longer contain "E" */
        assert!(! clause.contains(&term!["E"]));
    }
}
<|endoftext|>


<commit_before>use std::str::FromStr;
use std::error::Error;
use std::fmt;

use std::marker::PhantomData;

/// A data table consisting of varying column types and headers.
pub struct DataTable {
    pub data_cols: Vec<DataColumn>,
}

impl DataTable {
    pub fn empty() -> DataTable {
        DataTable { data_cols: Vec::new() }
    }

    pub fn cols(&self) -> usize {
        self.data_cols.len()
    }

    pub fn rows(&self) -> usize {
        if self.data_cols.len() > 0 {
            return self.data_cols[0].len();
        }

        0usize
    }

    /// Shrinks the table and it's underlying
    pub fn shrink_to_fit(&mut self) {
        for col in self.data_cols.iter_mut() {
            col.shrink_to_fit();
        }

        self.data_cols.shrink_to_fit();
    }

    /// Attempt to convert the data table into a single Vec.
    ///
    /// Uses column major ordering.
    ///
    /// # Failures
    ///
    /// - DataCastError : Returned then the data cannot be cast into the requested type.
    pub fn into_consistent_data<T: FromStr>(self) -> Result<Vec<T>, DataError> {
        let mut table_data = Vec::with_capacity(self.cols() * self.rows());
        for d in self.data_cols.into_iter() {
            match d.into_vec() {
                Ok(x) => table_data.extend(x),
                Err(e) => return Err(e),
            }
        }

        Ok(table_data)
    }
}

/// An error for attempting an into conversion.
#[derive(Debug)]
pub enum DataError {
    DataCastError,
}

impl fmt::Display for DataError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "DataCastError")
    }
}

impl Error for DataError {
    fn description(&self) -> &str {
        "Failed to cast data."
    }
}

/// A data column with a consistent data type.
pub struct DataColumn {
    data: Vec<String>,
}

impl DataColumn {
    pub fn empty() -> DataColumn {
        DataColumn {
            data: Vec::new(),
        }
    }

    pub fn len(&self) -> usize {
        self.data.len()
    }

    pub fn data(&self) -> &Vec<String> {
        &self.data
    }

    pub fn push(&mut self, val: &str) {
        self.data.push(val.to_owned());
    }

    pub fn into_vec<T: FromStr>(self) -> Result<Vec<T>, DataError> {
        let mut casted_data = Vec::<T>::with_capacity(self.data.len());

        for d in self.data.into_iter() {
            match T::from_str(&d[..]) {
                Ok(x) => casted_data.push(x),
                Err(_) => return Err(DataError::DataCastError),
            }
        }

        Ok(casted_data)
    }

    pub fn cast<T: FromStr>(&self) -> Option<Vec<T>> {
        let mut casted_data = Vec::<T>::with_capacity(self.data.len());

        for d in self.data.iter() {
            match T::from_str(&d[..]) {
                Ok(x) => casted_data.push(x),
                Err(_) => return None,
            }
        }

        Some(casted_data)
    }

    pub fn shrink_to_fit(&mut self) {
        self.data.shrink_to_fit();
    }

    pub fn into_iter<T: FromStr>(self) -> IntoColumnIterator<T> {
        IntoColumnIterator {
            column_data: self.data,
            index: 0usize,
            phantom: PhantomData::<T>,
        }
    }
}

pub struct IntoColumnIterator<T : FromStr> {
    column_data: Vec<String>,
    index: usize,
    phantom: PhantomData<T>,
}

impl<T : FromStr> Iterator for IntoColumnIterator<T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        match self.column_data.get(self.index) {
            None => None,
            Some(elt) => {
                self.index += 1;
                match T::from_str(&elt[..]) {
                    Ok(x) => Some(x),
                    Err(_) => None,
                }
            }
        }
    }
}<commit_msg>Implementing iterator for data columns.<commit_after>use std::str::FromStr;
use std::error::Error;
use std::fmt;
use std::vec::IntoIter;
use std;

/// A data table consisting of varying column types and headers.
pub struct DataTable {
    pub data_cols: Vec<DataColumn>,
}

impl DataTable {
    pub fn empty() -> DataTable {
        DataTable { data_cols: Vec::new() }
    }

    pub fn cols(&self) -> usize {
        self.data_cols.len()
    }

    pub fn rows(&self) -> usize {
        if self.data_cols.len() > 0 {
            return self.data_cols[0].len();
        }

        0usize
    }

    /// Shrinks the table and it's underlying
    pub fn shrink_to_fit(&mut self) {
        for col in self.data_cols.iter_mut() {
            col.shrink_to_fit();
        }

        self.data_cols.shrink_to_fit();
    }

    /// Attempt to convert the data table into a single Vec.
    ///
    /// Uses column major ordering.
    ///
    /// # Failures
    ///
    /// - DataCastError : Returned then the data cannot be cast into the requested type.
    pub fn into_consistent_data<T: FromStr>(self) -> Result<Vec<T>, DataError> {
        let mut table_data = Vec::with_capacity(self.cols() * self.rows());
        for d in self.data_cols.into_iter() {
            match d.into_vec() {
                Ok(x) => table_data.extend(x),
                Err(e) => return Err(e),
            }
        }

        Ok(table_data)
    }
}

/// An error for attempting an into conversion.
#[derive(Debug)]
pub enum DataError {
    DataCastError,
}

impl fmt::Display for DataError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "DataCastError")
    }
}

impl Error for DataError {
    fn description(&self) -> &str {
        "Failed to cast data."
    }
}

/// A data column with a consistent data type.
pub struct DataColumn {
    data: Vec<String>,
}

impl DataColumn {
    pub fn empty() -> DataColumn {
        DataColumn { data: Vec::new() }
    }

    pub fn len(&self) -> usize {
        self.data.len()
    }

    pub fn data(&self) -> &Vec<String> {
        &self.data
    }

    pub fn push(&mut self, val: &str) {
        self.data.push(val.to_owned());
    }

    pub fn into_vec<T: FromStr>(self) -> Result<Vec<T>, DataError> {
        let mut casted_data = Vec::<T>::with_capacity(self.data.len());

        for d in self.data.into_iter() {
            match T::from_str(&d[..]) {
                Ok(x) => casted_data.push(x),
                Err(_) => return Err(DataError::DataCastError),
            }
        }

        Ok(casted_data)
    }

    pub fn cast<T: FromStr>(&self) -> Option<Vec<T>> {
        let mut casted_data = Vec::<T>::with_capacity(self.data.len());

        for d in self.data.iter() {
            match T::from_str(&d[..]) {
                Ok(x) => casted_data.push(x),
                Err(_) => return None,
            }
        }

        Some(casted_data)
    }

    pub fn shrink_to_fit(&mut self) {
        self.data.shrink_to_fit();
    }

    /// Consumes self and returns an iterator which parses
    /// the data to the specified type returning results.
    ///
    /// The iterator will return a result on `next()` detailing
    /// the outcome of the parse.
    pub fn into_iter_cast<U: FromStr>
        (self)
         -> std::iter::Map<IntoIter<String>, fn(String) -> Result<U, <U as FromStr>::Err>>
        where U: FromStr
    {
        from_str_iter::<_, U>(self.data.into_iter())
    }
}

fn from_str_iter<I, U>
    (iter: I)
     -> std::iter::Map<I, fn(<I as Iterator>::Item) -> Result<U, <U as FromStr>::Err>>
    where I: Iterator,
          <I as Iterator>::Item: AsRef<str>,
          U: FromStr
{
    fn from_str_fn<T, U>(item: T) -> Result<U, <U as FromStr>::Err>
        where T: AsRef<str>,
              U: FromStr
    {
        FromStr::from_str(item.as_ref())
    }
    iter.map(from_str_fn)
}
<|endoftext|>