<commit_before>//
// Copyright 2020 The Project Oak Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

/// This module defines a trait and an implementation for representing the [`Runtime`] as
/// a Graphviz dot graph. This is used in the introspection service.
use crate::{ChannelHalf, ChannelHalfDirection, NodeId, Runtime};
use itertools::Itertools;
use log::info;
use std::{collections::HashSet, fmt::Write, string::String};

/// Trait that gives an identifier for a data structure that is suitable for
/// use with Graphviz/Dot.
pub trait DotIdentifier {
    fn dot_id(&self) -> String;
}

/// Trait that returns the path at which the debug introspection server will
/// show a page for a data structure.
pub trait HtmlPath {
    fn html_path(&self) -> String;
}

impl DotIdentifier for NodeId {
    fn dot_id(&self) -> String {
        format!("node{}", self.0)
    }
}

impl HtmlPath for NodeId {
    fn html_path(&self) -> String {
        format!("/node/{}", self.0)
    }
}

impl DotIdentifier for (NodeId, oak_abi::Handle) {
    fn dot_id(&self) -> String {
        format!("{}_{}", self.0.dot_id(), self.1)
    }
}

impl HtmlPath for (NodeId, oak_abi::Handle) {
    fn html_path(&self) -> String {
        format!("{}/{}", self.0.html_path(), self.1)
    }
}

// Graph-related methods for the Runtime.
impl Runtime {
    /// Generate a Graphviz dot graph that shows the current shape of the Nodes and Channels in
    /// the `Runtime`.
    pub fn graph(&self) -> String {
        let mut s = String::new();
        writeln!(&mut s, "digraph Runtime {{").unwrap();
        // Graph nodes for Oak Nodes and ABI handles.
        {
            writeln!(&mut s, "  {{").unwrap();
            writeln!(
                &mut s,
                "    node [shape=box style=filled fillcolor=red fontsize=24]"
            )
            .unwrap();
            let node_infos = self.node_infos.read().unwrap();
            for node_id in node_infos.keys().sorted() {
                let node_info = node_infos.get(node_id).unwrap();
                write!(
                    &mut s,
                    r###"    {} [label="{}" URL="{}"]"###,
                    node_id.dot_id(),
                    node_info.name,
                    node_id.html_path(),
                )
                .unwrap();
                if node_info.node_stopper.is_none() {
                    write!(&mut s, " [style=dashed]").unwrap();
                }
                writeln!(&mut s).unwrap();
            }
            writeln!(&mut s, "  }}").unwrap();
        }
        {
            writeln!(&mut s, "  {{").unwrap();
            writeln!(
                &mut s,
                "    node [shape=hexagon style=filled fillcolor=orange]"
            )
            .unwrap();
            let node_infos = self.node_infos.read().unwrap();
            for node_id in node_infos.keys().sorted() {
                let node_info = node_infos.get(node_id).unwrap();
                for handle in node_info.abi_handles.keys() {
                    writeln!(
                        &mut s,
                        r###"    {} [label="{}:{}" URL="{}"]"###,
                        (*node_id, *handle).dot_id(),
                        node_id.0,
                        handle,
                        (*node_id, *handle).html_path(),
                    )
                    .unwrap();
                }
            }
            writeln!(&mut s, "  }}").unwrap();
        }
        {
            writeln!(&mut s, "  {{").unwrap();
            writeln!(
                &mut s,
                "    node [shape=ellipse style=filled fillcolor=green]"
            )
            .unwrap();
            let node_infos = self.node_infos.read().unwrap();
            for node_id in node_infos.keys().sorted() {
                let node_info = node_infos.get(node_id).unwrap();
                for half in node_info.abi_handles.values() {
                    writeln!(
                        &mut s,
                        r###"    {} [label="{}"]"###,
                        half.dot_id(),
                        half.get_channel_name(),
                    )
                    .unwrap();
                }
            }
            writeln!(&mut s, "  }}").unwrap();
        }

        // Edges for connections between Nodes and channels and messages.
        let mut msg_counter = 0;
        {
            let node_infos = self.node_infos.read().unwrap();
            for node_id in node_infos.keys().sorted() {
                let node_info = node_infos.get(node_id).unwrap();
                for (handle, half) in &node_info.abi_handles {
                    match half.direction {
                        ChannelHalfDirection::Write => {
                            writeln!(
                                &mut s,
                                "  {} -> {} -> {}",
                                node_id.dot_id(),
                                (*node_id, *handle).dot_id(),
                                half.dot_id()
                            )
                            .unwrap();
                        }
                        ChannelHalfDirection::Read => {
                            writeln!(
                                &mut s,
                                "  {} -> {} -> {}",
                                half.dot_id(),
                                (*node_id, *handle).dot_id(),
                                node_id.dot_id()
                            )
                            .unwrap();
                        }
                    }
                    // Include messages in the channel.
                    let messages = half.get_messages();
                    if messages.len() > 0 {
                        writeln!(&mut s, "  {{").unwrap();
                        writeln!(
                            &mut s,
                            r###"    node [shape=rect fontsize=10 label="msg"]"###
                        )
                        .unwrap();
                        // Messages have no identifier, so just use a count (and don't make it
                        // visible to the user).
                        let mut prev_graph_node = half.dot_id();
                        for msg in messages.iter() {
                            msg_counter += 1;
                            let graph_node = format!("msg{}", msg_counter);
                            writeln!(
                                &mut s,
                                "    {} -> {} [style=dashed arrowhead=none]",
                                graph_node, prev_graph_node
                            )
                            .unwrap();
                            for half in &msg.channels {
                                match half.direction {
                                    ChannelHalfDirection::Write => {
                                        writeln!(&mut s, "    {} -> {}", graph_node, half.dot_id())
                                            .unwrap();
                                    }
                                    ChannelHalfDirection::Read => {
                                        writeln!(&mut s, "    {} -> {}", half.dot_id(), graph_node)
                                            .unwrap();
                                    }
                                }
                            }
                            prev_graph_node = graph_node;
                        }
                        writeln!(&mut s, "  }}").unwrap();
                    }
                }
            }
        }

        writeln!(&mut s, "}}").unwrap();
        s
    }

    /// Generate an HTML page that describes the internal state of the [`Runtime`].
    pub fn html(&self) -> String {
        let mut s = String::new();
        writeln!(&mut s, "
Nodes
").unwrap();
        writeln!(&mut s, r###"
Show as graph
"###).unwrap();
        {
            let node_infos = self.node_infos.read().unwrap();
            for node_id in node_infos.keys().sorted() {
                let node_info = node_infos.get(node_id).unwrap();
                write!(
                    &mut s,
                    r###"
{:?} => <tt>{:?}</tt>"###,
                    node_id.html_path(),
                    node_id,
                    node_info
                )
                .unwrap();
            }
        }
        writeln!(&mut s, "
").unwrap();
        s
    }

    /// Generate an HTML page that includes the current counts of Nodes and channels.
    /// May be slow to generate, as it involves exploring reachable channels recursively.
    pub(crate) fn html_counts(&self) -> String {
        let (node_count, channel_count) = self.object_counts();
        format!(
            "
Object Counts: Nodes={} Channels={}",
            node_count, channel_count,
        )
    }

    /// Return counts of the number of Nodes and channels currently in existence.
    /// May be slow to generate, as it involves exploring reachable channels recursively.
    pub fn object_counts(&self) -> (u32, u32) {
        let mut node_count = 0;
        let mut channel_ids = HashSet::<String>::new();
        let mut visitor = |half: &ChannelHalf| {
            let channel_id = half.dot_id();
            if channel_ids.contains(&channel_id) {
                false
            } else {
                channel_ids.insert(channel_id);
                // Not seen this ChannelId before, so need to visit its children.
                true
            }
        };
        {
            let node_infos = self.node_infos.read().unwrap();
            for node_id in node_infos.keys().sorted() {
                node_count += 1;
                let node_info = node_infos.get(node_id).unwrap();
                for half in node_info.abi_handles.values() {
                    half.visit_halves(&mut visitor);
                }
            }
        }
        info!(
            "Counted {} nodes and {} channels",
            node_count,
            channel_ids.len()
        );
        (node_count as u32, channel_ids.len() as u32)
    }

    /// Generate an HTML page that describes the internal state of a specific Node.
    pub(crate) fn html_for_node(&self, id: u64) -> Option<String> {
        let node_id = NodeId(id);
        let node_infos = self.node_infos.read().unwrap();
        let node_info = node_infos.get(&node_id)?;
        let mut s = String::new();
        write!(&mut s, "
{}
", node_info.name).unwrap();
        if let Some(node_stopper) = &node_info.node_stopper {
            write!(
                &mut s,
                "
Node thread is currently running, join_handle={:?}",
                node_stopper.join_handle
            )
            .unwrap();
        } else {
            write!(&mut s, "
No current thread for Node.").unwrap();
        }
        write!(&mut s, "
Label={:?}
Handles:
", node_info.label).unwrap();
        for (handle, half) in &node_info.abi_handles {
            write!(
                &mut s,
                r###"
Handle {} => {:?}"###,
                (node_id, *handle).html_path(),
                handle,
                half
            )
            .unwrap();
        }
        Some(s)
    }

    /// Generate an HTML page that describes the channel accessible via an ABI handle for the
    /// specified Node.
    pub(crate) fn html_for_handle(&self, node_id: u64, handle: oak_abi::Handle) -> Option<String> {
        let node_id = NodeId(node_id);
        let node_infos = self.node_infos.read().unwrap();
        let node_info = node_infos.get(&node_id)?;
        let half = node_info.abi_handles.get(&handle)?;
        let mut s = String::new();
        write!(
            &mut s,
            r###"
Node {} Handle {}
"###,
            node_id.html_path(),
            node_info.name,
            handle,
        )
        .unwrap();
        write!(&mut s, r###"
Maps to {:?}"###, half).unwrap();
        Some(s)
    }
}
<commit_msg>runtime: de-dup channels in introspection graph (#1719)<commit_after>//
// Copyright 2020 The Project Oak Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

/// This module defines a trait and an implementation for representing the [`Runtime`] as
/// a Graphviz dot graph. This is used in the introspection service.
use crate::{ChannelHalf, ChannelHalfDirection, NodeId, Runtime};
use itertools::Itertools;
use log::info;
use std::{collections::HashSet, fmt::Write, string::String};

/// Trait that gives an identifier for a data structure that is suitable for
/// use with Graphviz/Dot.
pub trait DotIdentifier {
    fn dot_id(&self) -> String;
}

/// Trait that returns the path at which the debug introspection server will
/// show a page for a data structure.
pub trait HtmlPath {
    fn html_path(&self) -> String;
}

impl DotIdentifier for NodeId {
    fn dot_id(&self) -> String {
        format!("node{}", self.0)
    }
}

impl HtmlPath for NodeId {
    fn html_path(&self) -> String {
        format!("/node/{}", self.0)
    }
}

impl DotIdentifier for (NodeId, oak_abi::Handle) {
    fn dot_id(&self) -> String {
        format!("{}_{}", self.0.dot_id(), self.1)
    }
}

impl HtmlPath for (NodeId, oak_abi::Handle) {
    fn html_path(&self) -> String {
        format!("{}/{}", self.0.html_path(), self.1)
    }
}

// Graph-related methods for the Runtime.
impl Runtime {
    /// Generate a Graphviz dot graph that shows the current shape of the Nodes and Channels in
    /// the `Runtime`.
    pub fn graph(&self) -> String {
        let mut s = String::new();
        writeln!(&mut s, "digraph Runtime {{").unwrap();
        // Graph nodes for Oak Nodes and ABI handles.
        {
            writeln!(&mut s, "  {{").unwrap();
            writeln!(
                &mut s,
                "    node [shape=box style=filled fillcolor=red fontsize=24]"
            )
            .unwrap();
            let node_infos = self.node_infos.read().unwrap();
            for node_id in node_infos.keys().sorted() {
                let node_info = node_infos.get(node_id).unwrap();
                write!(
                    &mut s,
                    r###"    {} [label="{}" URL="{}"]"###,
                    node_id.dot_id(),
                    node_info.name,
                    node_id.html_path(),
                )
                .unwrap();
                if node_info.node_stopper.is_none() {
                    write!(&mut s, " [style=dashed]").unwrap();
                }
                writeln!(&mut s).unwrap();
            }
            writeln!(&mut s, "  }}").unwrap();
        }
        {
            writeln!(&mut s, "  {{").unwrap();
            writeln!(
                &mut s,
                "    node [shape=hexagon style=filled fillcolor=orange]"
            )
            .unwrap();
            let node_infos = self.node_infos.read().unwrap();
            for node_id in node_infos.keys().sorted() {
                let node_info = node_infos.get(node_id).unwrap();
                for handle in node_info.abi_handles.keys() {
                    writeln!(
                        &mut s,
                        r###"    {} [label="{}:{}" URL="{}"]"###,
                        (*node_id, *handle).dot_id(),
                        node_id.0,
                        handle,
                        (*node_id, *handle).html_path(),
                    )
                    .unwrap();
                }
            }
            writeln!(&mut s, "  }}").unwrap();
        }
        {
            writeln!(&mut s, "  {{").unwrap();
            writeln!(
                &mut s,
                "    node [shape=ellipse style=filled fillcolor=green]"
            )
            .unwrap();
            let mut seen = HashSet::new();
            let node_infos = self.node_infos.read().unwrap();
            for node_id in node_infos.keys().sorted() {
                let node_info = node_infos.get(node_id).unwrap();
                for half in node_info.abi_handles.values() {
                    if seen.insert(half.dot_id()) {
                        writeln!(
                            &mut s,
                            r###"    {} [label="{}"]"###,
                            half.dot_id(),
                            half.get_channel_name(),
                        )
                        .unwrap();
                    }
                }
            }
            writeln!(&mut s, "  }}").unwrap();
        }

        // Edges for connections between Nodes and channels and messages.
        let mut msg_counter = 0;
        {
            let node_infos = self.node_infos.read().unwrap();
            for node_id in node_infos.keys().sorted() {
                let node_info = node_infos.get(node_id).unwrap();
                for (handle, half) in &node_info.abi_handles {
                    match half.direction {
                        ChannelHalfDirection::Write => {
                            writeln!(
                                &mut s,
                                "  {} -> {} -> {}",
                                node_id.dot_id(),
                                (*node_id, *handle).dot_id(),
                                half.dot_id()
                            )
                            .unwrap();
                        }
                        ChannelHalfDirection::Read => {
                            writeln!(
                                &mut s,
                                "  {} -> {} -> {}",
                                half.dot_id(),
                                (*node_id, *handle).dot_id(),
                                node_id.dot_id()
                            )
                            .unwrap();
                        }
                    }
                    // Include messages in the channel.
                    let messages = half.get_messages();
                    if messages.len() > 0 {
                        writeln!(&mut s, "  {{").unwrap();
                        writeln!(
                            &mut s,
                            r###"    node [shape=rect fontsize=10 label="msg"]"###
                        )
                        .unwrap();
                        // Messages have no identifier, so just use a count (and don't make it
                        // visible to the user).
                        let mut prev_graph_node = half.dot_id();
                        for msg in messages.iter() {
                            msg_counter += 1;
                            let graph_node = format!("msg{}", msg_counter);
                            writeln!(
                                &mut s,
                                "    {} -> {} [style=dashed arrowhead=none]",
                                graph_node, prev_graph_node
                            )
                            .unwrap();
                            for half in &msg.channels {
                                match half.direction {
                                    ChannelHalfDirection::Write => {
                                        writeln!(&mut s, "    {} -> {}", graph_node, half.dot_id())
                                            .unwrap();
                                    }
                                    ChannelHalfDirection::Read => {
                                        writeln!(&mut s, "    {} -> {}", half.dot_id(), graph_node)
                                            .unwrap();
                                    }
                                }
                            }
                            prev_graph_node = graph_node;
                        }
                        writeln!(&mut s, "  }}").unwrap();
                    }
                }
            }
        }

        writeln!(&mut s, "}}").unwrap();
        s
    }

    /// Generate an HTML page that describes the internal state of the [`Runtime`].
    pub fn html(&self) -> String {
        let mut s = String::new();
        writeln!(&mut s, "
Nodes
").unwrap();
        writeln!(&mut s, r###"
Show as graph
"###).unwrap();
        {
            let node_infos = self.node_infos.read().unwrap();
            for node_id in node_infos.keys().sorted() {
                let node_info = node_infos.get(node_id).unwrap();
                write!(
                    &mut s,
                    r###"
{:?} => <tt>{:?}</tt>"###,
                    node_id.html_path(),
                    node_id,
                    node_info
                )
                .unwrap();
            }
        }
        writeln!(&mut s, "
").unwrap();
        s
    }

    /// Generate an HTML page that includes the current counts of Nodes and channels.
    /// May be slow to generate, as it involves exploring reachable channels recursively.
    pub(crate) fn html_counts(&self) -> String {
        let (node_count, channel_count) = self.object_counts();
        format!(
            "
Object Counts: Nodes={} Channels={}",
            node_count, channel_count,
        )
    }

    /// Return counts of the number of Nodes and channels currently in existence.
    /// May be slow to generate, as it involves exploring reachable channels recursively.
    pub fn object_counts(&self) -> (u32, u32) {
        let mut node_count = 0;
        let mut channel_ids = HashSet::<String>::new();
        let mut visitor = |half: &ChannelHalf| {
            let channel_id = half.dot_id();
            if channel_ids.contains(&channel_id) {
                false
            } else {
                channel_ids.insert(channel_id);
                // Not seen this ChannelId before, so need to visit its children.
                true
            }
        };
        {
            let node_infos = self.node_infos.read().unwrap();
            for node_id in node_infos.keys().sorted() {
                node_count += 1;
                let node_info = node_infos.get(node_id).unwrap();
                for half in node_info.abi_handles.values() {
                    half.visit_halves(&mut visitor);
                }
            }
        }
        info!(
            "Counted {} nodes and {} channels",
            node_count,
            channel_ids.len()
        );
        (node_count as u32, channel_ids.len() as u32)
    }

    /// Generate an HTML page that describes the internal state of a specific Node.
    pub(crate) fn html_for_node(&self, id: u64) -> Option<String> {
        let node_id = NodeId(id);
        let node_infos = self.node_infos.read().unwrap();
        let node_info = node_infos.get(&node_id)?;
        let mut s = String::new();
        write!(&mut s, "
{}
", node_info.name).unwrap();
        if let Some(node_stopper) = &node_info.node_stopper {
            write!(
                &mut s,
                "
Node thread is currently running, join_handle={:?}",
                node_stopper.join_handle
            )
            .unwrap();
        } else {
            write!(&mut s, "
No current thread for Node.").unwrap();
        }
        write!(&mut s, "
Label={:?}
Handles:
", node_info.label).unwrap();
        for (handle, half) in &node_info.abi_handles {
            write!(
                &mut s,
                r###"
Handle {} => {:?}"###,
                (node_id, *handle).html_path(),
                handle,
                half
            )
            .unwrap();
        }
        Some(s)
    }

    /// Generate an HTML page that describes the channel accessible via an ABI handle for the
    /// specified Node.
    pub(crate) fn html_for_handle(&self, node_id: u64, handle: oak_abi::Handle) -> Option<String> {
        let node_id = NodeId(node_id);
        let node_infos = self.node_infos.read().unwrap();
        let node_info = node_infos.get(&node_id)?;
        let half = node_info.abi_handles.get(&handle)?;
        let mut s = String::new();
        write!(
            &mut s,
            r###"
Node {} Handle {}
"###,
            node_id.html_path(),
            node_info.name,
            handle,
        )
        .unwrap();
        write!(&mut s, r###"
Maps to {:?}"###, half).unwrap();
        Some(s)
    }
}
<|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
use super::model;
use regex;
use rs_es;
use rs_es::query::Query as rs_q;
use rs_es::operations::search::SearchResult;
use rs_es::units as rs_u;
use mimir;
use serde_json;
use serde;
use mimir::objects::{MimirObject, Admin, Addr, Stop};

fn build_rs_client(cnx: &String) -> rs_es::Client {
    let re = regex::Regex::new(r"(?:https?://)?(?P<host>.+?):(?P<port>\d+)").unwrap();
    let cap = re.captures(&cnx).unwrap();
    let host = cap.name("host").unwrap().as_str();
    let port = cap.name("port").unwrap().as_str().parse::<u32>().unwrap();

    rs_es::Client::new(&host, port)
}

/// takes a ES json blob and build a Place from it
/// it uses the _type field of ES to know which type of the Place enum to fill
pub fn make_place(doc_type: String, value: Option<Box<serde_json::Value>>) -> Option<mimir::Place> {
    value.and_then(|v| {
        fn convert<T: serde::Deserialize>(v: serde_json::Value,
                                          f: fn(T) -> mimir::Place)
                                          -> Option<mimir::Place> {
            serde_json::from_value::<T>(v)
                .map_err(|err| warn!("Impossible to load ES result: {}", err))
                .ok()
                .map(f)
        }
        match doc_type.as_ref() {
            "addr" => convert(*v, mimir::Place::Addr),
            "street" => convert(*v, mimir::Place::Street),
            "admin" => convert(*v, mimir::Place::Admin),
            "poi" => convert(*v, mimir::Place::Poi),
            "stop" => convert(*v, mimir::Place::Stop),
            _ => {
                warn!("unknown ES return value, _type field = {}", doc_type);
                None
            }
        }
    })
}

#[derive(Debug, Eq, PartialEq)]
enum MatchType {
    Prefix,
    Fuzzy,
}

fn build_query(q: &str,
               match_type: MatchType,
               coord: &Option<model::Coord>,
               shape: Option<Vec<rs_es::units::Location>>)
               -> rs_es::query::Query {
    use rs_es::query::functions::Function;
    // we order the type of object we want
    // Note: the addresses are boosted more because even if we don't want them first
    // because they are more severely filtered
    let boost_addr = rs_q::build_term("_type", Addr::doc_type()).with_boost(5000).build();
    let boost_admin = rs_q::build_term("_type", Admin::doc_type()).with_boost(3000).build();
    let boost_stop = rs_q::build_term("_type", Stop::doc_type()).with_boost(2000).build();

    let main_match_type = match match_type {
        MatchType::Prefix => "label.prefix",
        MatchType::Fuzzy => "label.ngram",
    };

    let boost_main_match_query = rs_q::build_match(main_match_type.to_string(), q.to_string())
        .with_boost(500)
        .build();

    let boost_zipcode_match_query = rs_q::build_match("zip_codes.prefix".to_string(),
                                                      q.to_string())
        .with_boost(100)
        .build();

    let mut should_query = vec![boost_addr,
                                boost_admin,
                                boost_stop,
                                boost_main_match_query,
                                boost_zipcode_match_query];

    // for fuzzy search we also search by the prefix index (with a greater boost than ngram)
    // to have better results
    if match_type == MatchType::Fuzzy {
        should_query.push(rs_q::build_match("label.prefix", q.to_string())
            .with_boost(1000)
            .build());
    }

    if let &Some(ref c) = coord {
        // if we have coordinate, we boost we result near this coordinate
        let boost_on_proximity =
            rs_q::build_function_score()
                .with_boost_mode(rs_es::query::compound::BoostMode::Multiply)
                .with_boost(1500)
                .with_function(Function::build_decay("coord",
                                           rs_u::Location::LatLon(c.lat, c.lon),
                                           rs_u::Distance::new(50f64,
                                                               rs_u::DistanceUnit::Kilometer))
                                   .build_exp())
                .build();
        should_query.push(boost_on_proximity);
    } else {
        // if we don't have coords, we take the field `weight` into account
        let boost_on_weight = rs_q::build_function_score()
            .with_boost_mode(rs_es::query::compound::BoostMode::Multiply)
            .with_boost(500)
            .with_query(rs_q::build_match_all().build())
            .with_function(Function::build_field_value_factor("weight")
                .with_factor(1)
                .with_modifier(rs_es::query::functions::Modifier::Log1p)
                .build())
            .build();
        should_query.push(boost_on_weight);
    }

    let sub_query = rs_q::build_bool()
        .with_should(should_query)
        .build();

    // filter to handle house number
    // we either want:
    // * to exactly match the document house_number
    // * or that the document has no house_number
    let first_condition = rs_q::build_bool()
        .with_should(vec![rs_q::build_bool()
                              .with_must_not(rs_q::build_exists("house_number").build())
                              .build(),
                          rs_q::build_match("house_number", q.to_string()).build()])
        .build();

    use rs_es::query::MinimumShouldMatch;

    let second_condition = match match_type {
        MatchType::Prefix => {
            rs_q::build_multi_match(vec!["label.prefix".to_string(),
                                         "zip_codes.prefix".to_string()],
                                    q.to_string())
                .with_type(rs_es::query::full_text::MatchQueryType::CrossFields)
                .with_operator("and")
                .build()
        }
        // for fuzzy search we lower our expectation and we accept 50% of token match
        MatchType::Fuzzy => {
            rs_q::build_multi_match(vec!["label.prefix".to_string(),
                                         "zip_codes.prefix".to_string()],
                                    q.to_string())
                .with_minimum_should_match(MinimumShouldMatch::from(40f64))
                .build()
        }
    };

    let mut must = vec![first_condition, second_condition];

    if let Some(s) = shape {
        must.push(rs_q::build_geo_polygon("coord", s).build());
    }
    let filter = rs_q::build_bool()
        .with_must(must)
        .build();

    rs_q::build_bool()
        .with_must(vec![sub_query])
        .with_filter(filter)
        .build()
}

fn is_existing_index(client: &mut rs_es::Client, index: &str) -> bool {
    !index.is_empty() && client.open_index(&index).is_ok()
}

fn get_indexes_by_type(a_type: &str) -> String {
    let doc_type = match a_type {
        "public_transport:stop_area" => "stop",
        "city" => "admin",
        "house" => "addr",
        _ => a_type,
    };

    format!("munin_{}", doc_type)
}

fn make_indexes_impl<F: FnMut(&str) -> bool>(all_data: bool,
                                             pt_dataset_index: &Option<String>,
                                             types: &Option<Vec<&str>>,
                                             mut is_existing_index: F)
                                             -> Vec<String> {
    if all_data {
        return vec!["munin".to_string()];
    }

    let mut result: Vec<String> = vec![];
    let mut push = |result: &mut Vec<_>, i: &str| if is_existing_index(i) {
        result.push(i.into());
    };
    if let Some(ref types) = *types {
        for type_ in types.iter().filter(|t| **t != "public_transport:stop_area") {
            push(&mut result, &get_indexes_by_type(type_));
        }
        match *pt_dataset_index {
            Some(ref index) if types.contains(&"public_transport:stop_area") => {
                push(&mut result, index)
            }
            _ => (),
        }
    } else {
        push(&mut result, &"munin_geo_data".to_string());

        if let Some(ref dataset) = *pt_dataset_index {
            push(&mut result, &dataset.clone());
        }
    }

    result
}

fn make_indexes(all_data: bool,
                pt_dataset_index: &Option<String>,
                types: &Option<Vec<&str>>,
                client: &mut rs_es::Client)
                -> Vec<String> {
    make_indexes_impl(all_data,
                      pt_dataset_index,
                      types,
                      |index| is_existing_index(client, index))
}

fn query(q: &str,
         pt_dataset: &Option<&str>,
         all_data: bool,
         cnx: &str,
         match_type: MatchType,
         offset: u64,
         limit: u64,
         coord: &Option<model::Coord>,
         shape: Option<Vec<rs_es::units::Location>>,
         types: &Option<Vec<&str>>)
         -> Result<Vec<mimir::Place>, rs_es::error::EsError> {
    let query = build_query(q, match_type, coord, shape);

    let mut client = build_rs_client(&cnx.to_string());

    let pt_dataset_index = pt_dataset.map(|d| format!("munin_stop_{}", d));
    let indexes = make_indexes(all_data, &pt_dataset_index, types, &mut client);

    debug!("ES indexes: {:?}", indexes);

    if indexes.is_empty() {
        // if there is no indexes, rs_es search with index "_all"
        // but we want to return emtpy response in this case.
        return Ok(vec![]);
    }

    let result: SearchResult<serde_json::Value> = try!(client.search_query()
        .with_indexes(&indexes.iter().map(|index| index.as_str()).collect::<Vec<&str>>())
        .with_query(&query)
        .with_from(offset)
        .with_size(limit)
        .send());

    debug!("{} documents found", result.hits.total);

    // for the moment rs-es does not handle enum Document,
    // so we need to convert the ES glob to a Place
    Ok(result.hits
        .hits
        .into_iter()
        .filter_map(|hit| make_place(hit.doc_type, hit.source))
        .collect())
}

pub fn autocomplete(q: &str,
                    pt_dataset: &Option<&str>,
                    all_data: bool,
                    offset: u64,
                    limit: u64,
                    coord: Option<model::Coord>,
                    cnx: &str,
                    shape: Option<Vec<(f64, f64)>>,
                    types: Option<Vec<&str>>)
                    -> Result<Vec<mimir::Place>, rs_es::error::EsError> {
    fn make_shape(shape: &Option<Vec<(f64, f64)>>) -> Option<Vec<rs_es::units::Location>> {
        shape.as_ref().map(|v| v.iter().map(|&l| l.into()).collect())
    }

    // First we try a prety exact match on the prefix.
    // If there are no results then we do a new fuzzy search (matching ngrams)
    let results = try!(query(&q,
                             &pt_dataset,
                             all_data,
                             cnx,
                             MatchType::Prefix,
                             offset,
                             limit,
                             &coord,
                             make_shape(&shape),
                             &types));
    if results.is_empty() {
        query(&q,
              &pt_dataset,
              all_data,
              cnx,
              MatchType::Fuzzy,
              offset,
              limit,
              &coord,
              make_shape(&shape),
              &types)
    } else {
        Ok(results)
    }
}

#[test]
fn test_make_indexes_impl() {
    // all_data
    assert_eq!(make_indexes_impl(true, &None, &None, |_index| true),
               vec!["munin"]);

    // no dataset and no types
    assert_eq!(make_indexes_impl(false, &None, &None, |_index| true),
               vec!["munin_geo_data"]);

    // dataset fr + no types
    assert_eq!(make_indexes_impl(false,
                                 &Some("munin_stop_fr".to_string()),
                                 &None,
                                 |_index| true),
               vec!["munin_geo_data", "munin_stop_fr"]);

    // no dataset + types poi, city, street, house and public_transport:stop_area
    // => munin_stop is not included
    assert_eq!(make_indexes_impl(false,
                                 &None,
                                 &Some(vec!["poi",
                                            "city",
                                            "street",
                                            "house",
                                            "public_transport:stop_area"]),
                                 |_index| true),
               vec!["munin_poi", "munin_admin", "munin_street", "munin_addr"]);

    // no dataset fr + type public_transport:stop_area only
    assert_eq!(make_indexes_impl(false,
                                 &None,
                                 &Some(vec!["public_transport:stop_area"]),
                                 |_index| true),
               Vec::<String>::new());

    // dataset fr + types poi, city, street, house and public_transport:stop_area
    assert_eq!(make_indexes_impl(false,
                                 &Some("munin_stop_fr".to_string()),
                                 &Some(vec!["poi",
                                            "city",
                                            "street",
                                            "house",
                                            "public_transport:stop_area"]),
                                 |_index| true),
               vec!["munin_poi", "munin_admin", "munin_street", "munin_addr", "munin_stop_fr"]);

    // dataset fr types poi, city, street, house without public_transport:stop_area
    //  => munin_stop_fr is not included
    assert_eq!(make_indexes_impl(false,
                                 &Some("munin_stop_fr".to_string()),
                                 &Some(vec!["poi", "city", "street", "house"]),
                                 |_index| true),
               vec!["munin_poi", "munin_admin", "munin_street", "munin_addr"]);
}
<commit_msg>add comments<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
use super::model;
use regex;
use rs_es;
use rs_es::query::Query as rs_q;
use rs_es::operations::search::SearchResult;
use rs_es::units as rs_u;
use mimir;
use serde_json;
use serde;
use mimir::objects::{MimirObject, Admin, Addr, Stop};

fn build_rs_client(cnx: &String) -> rs_es::Client {
    let re = regex::Regex::new(r"(?:https?://)?(?P<host>.+?):(?P<port>\d+)").unwrap();
    let cap = re.captures(&cnx).unwrap();
    let host = cap.name("host").unwrap().as_str();
    let port = cap.name("port").unwrap().as_str().parse::<u32>().unwrap();

    rs_es::Client::new(&host, port)
}

/// takes a ES json blob and build a Place from it
/// it uses the _type field of ES to know which type of the Place enum to fill
pub fn make_place(doc_type: String, value: Option<Box<serde_json::Value>>) -> Option<mimir::Place> {
    value.and_then(|v| {
        fn convert<T: serde::Deserialize>(v: serde_json::Value,
                                          f: fn(T) -> mimir::Place)
                                          -> Option<mimir::Place> {
            serde_json::from_value::<T>(v)
                .map_err(|err| warn!("Impossible to load ES result: {}", err))
                .ok()
                .map(f)
        }
        match doc_type.as_ref() {
            "addr" => convert(*v, mimir::Place::Addr),
            "street" => convert(*v, mimir::Place::Street),
            "admin" => convert(*v, mimir::Place::Admin),
            "poi" => convert(*v, mimir::Place::Poi),
            "stop" => convert(*v, mimir::Place::Stop),
            _ => {
                warn!("unknown ES return value, _type field = {}", doc_type);
                None
            }
        }
    })
}

#[derive(Debug, Eq, PartialEq)]
enum MatchType {
    Prefix,
    Fuzzy,
}

fn build_query(q: &str,
               match_type: MatchType,
               coord: &Option<model::Coord>,
               shape: Option<Vec<rs_es::units::Location>>)
               -> rs_es::query::Query {
    use rs_es::query::functions::Function;
    // we order the type of object we want
    // Note: the addresses are boosted more because even if we don't want them first
    // because they are more severely filtered
    let boost_addr = rs_q::build_term("_type", Addr::doc_type()).with_boost(5000).build();
    let boost_admin = rs_q::build_term("_type", Admin::doc_type()).with_boost(3000).build();
    let boost_stop = rs_q::build_term("_type", Stop::doc_type()).with_boost(2000).build();

    let main_match_type = match match_type {
        MatchType::Prefix => "label.prefix",
        MatchType::Fuzzy => "label.ngram",
    };

    let boost_main_match_query = rs_q::build_match(main_match_type.to_string(), q.to_string())
        .with_boost(500)
        .build();

    let boost_zipcode_match_query = rs_q::build_match("zip_codes.prefix".to_string(),
                                                      q.to_string())
        .with_boost(100)
        .build();

    let mut should_query = vec![boost_addr,
                                boost_admin,
                                boost_stop,
                                boost_main_match_query,
                                boost_zipcode_match_query];

    // for fuzzy search we also search by the prefix index (with a greater boost than ngram)
    // to have better results
    if match_type == MatchType::Fuzzy {
        should_query.push(rs_q::build_match("label.prefix", q.to_string())
            .with_boost(1000)
            .build());
    }

    if let &Some(ref c) = coord {
        // if we have coordinate, we boost we result near this coordinate
        let boost_on_proximity =
            rs_q::build_function_score()
                .with_boost_mode(rs_es::query::compound::BoostMode::Multiply)
                .with_boost(1500)
                .with_function(Function::build_decay("coord",
                                           rs_u::Location::LatLon(c.lat, c.lon),
                                           rs_u::Distance::new(50f64,
                                                               rs_u::DistanceUnit::Kilometer))
                                   .build_exp())
                .build();
        should_query.push(boost_on_proximity);
    } else {
        // if we don't have coords, we take the field `weight` into account
        let boost_on_weight = rs_q::build_function_score()
            .with_boost_mode(rs_es::query::compound::BoostMode::Multiply)
            .with_boost(500)
            .with_query(rs_q::build_match_all().build())
            .with_function(Function::build_field_value_factor("weight")
                .with_factor(1)
                .with_modifier(rs_es::query::functions::Modifier::Log1p)
                .build())
            .build();
        should_query.push(boost_on_weight);
    }

    let sub_query = rs_q::build_bool()
        .with_should(should_query)
        .build();

    // filter to handle house number
    // we either want:
    // * to exactly match the document house_number
    // * or that the document has no house_number
    let first_condition = rs_q::build_bool()
        .with_should(vec![rs_q::build_bool()
                              .with_must_not(rs_q::build_exists("house_number").build())
                              .build(),
                          rs_q::build_match("house_number", q.to_string()).build()])
        .build();

    use rs_es::query::MinimumShouldMatch;

    let second_condition = match match_type {
        // When the match type is Prefix, we want to use every possible information even though
        // these are not present in label, for instance, the zip_code. The cross_fields match type
        // allows to do the trick.

        // Ex:
        //   q : 20 rue hector malot 75012
        // WITHOUT the cross_fields match type, it will match neither "label" nor "zip_codes" and
        // the request will be treated by Fuzzy later, it's a pitty, because the adresse is actually
        // well spelt.
        // WITH the cross_fields match type, the request will be spilted into terms to match
        // "label" and "zip_codes"
        MatchType::Prefix => {
            rs_q::build_multi_match(vec!["label.prefix".to_string(),
                                         "zip_codes.prefix".to_string()],
                                    q.to_string())
                .with_type(rs_es::query::full_text::MatchQueryType::CrossFields)
                .with_operator("and")
                .build()
        }
        // for fuzzy search we lower our expectation and we accept 45% of token match
        MatchType::Fuzzy => {
            rs_q::build_multi_match(vec!["label.prefix".to_string(),
                                         "zip_codes.prefix".to_string()],
                                    q.to_string())
                .with_minimum_should_match(MinimumShouldMatch::from(45f64))
                .build()
        }
    };

    let mut must = vec![first_condition, second_condition];

    if let Some(s) = shape {
        must.push(rs_q::build_geo_polygon("coord", s).build());
    }
    let filter = rs_q::build_bool()
        .with_must(must)
        .build();

    rs_q::build_bool()
        .with_must(vec![sub_query])
        .with_filter(filter)
        .build()
}

fn is_existing_index(client: &mut rs_es::Client, index: &str) -> bool {
    !index.is_empty() && client.open_index(&index).is_ok()
}

fn get_indexes_by_type(a_type: &str) -> String {
    let doc_type = match a_type {
        "public_transport:stop_area" => "stop",
        "city" => "admin",
        "house" => "addr",
        _ => a_type,
    };

    format!("munin_{}", doc_type)
}

fn make_indexes_impl<F: FnMut(&str) -> bool>(all_data: bool,
                                             pt_dataset_index: &Option<String>,
                                             types: &Option<Vec<&str>>,
                                             mut is_existing_index: F)
                                             -> Vec<String> {
    if all_data {
        return vec!["munin".to_string()];
    }

    let mut result: Vec<String> = vec![];
    let mut push = |result: &mut Vec<_>, i: &str| if is_existing_index(i) {
        result.push(i.into());
    };
    if let Some(ref types) = *types {
        for type_ in types.iter().filter(|t| **t != "public_transport:stop_area") {
            push(&mut result, &get_indexes_by_type(type_));
        }
        match *pt_dataset_index {
            Some(ref index) if types.contains(&"public_transport:stop_area") => {
                push(&mut result, index)
            }
            _ => (),
        }
    } else {
        push(&mut result, &"munin_geo_data".to_string());

        if let Some(ref dataset) = *pt_dataset_index {
            push(&mut result, &dataset.clone());
        }
    }

    result
}

fn make_indexes(all_data: bool,
                pt_dataset_index: &Option<String>,
                types: &Option<Vec<&str>>,
                client: &mut rs_es::Client)
                -> Vec<String> {
    make_indexes_impl(all_data,
                      pt_dataset_index,
                      types,
                      |index| is_existing_index(client, index))
}

fn query(q: &str,
         pt_dataset: &Option<&str>,
         all_data: bool,
         cnx: &str,
         match_type: MatchType,
         offset: u64,
         limit: u64,
         coord: &Option<model::Coord>,
         shape: Option<Vec<rs_es::units::Location>>,
         types: &Option<Vec<&str>>)
         -> Result<Vec<mimir::Place>, rs_es::error::EsError> {
    let query = build_query(q, match_type, coord, shape);

    let mut client = build_rs_client(&cnx.to_string());

    let pt_dataset_index = pt_dataset.map(|d| format!("munin_stop_{}", d));
    let indexes = make_indexes(all_data, &pt_dataset_index, types, &mut client);

    debug!("ES indexes: {:?}", indexes);

    if indexes.is_empty() {
        // if there is no indexes, rs_es search with index "_all"
        // but we want to return emtpy response in this case.
        return Ok(vec![]);
    }

    let result: SearchResult<serde_json::Value> = try!(client.search_query()
        .with_indexes(&indexes.iter().map(|index| index.as_str()).collect::<Vec<&str>>())
        .with_query(&query)
        .with_from(offset)
        .with_size(limit)
        .send());

    debug!("{} documents found", result.hits.total);

    // for the moment rs-es does not handle enum Document,
    // so we need to convert the ES glob to a Place
    Ok(result.hits
        .hits
        .into_iter()
        .filter_map(|hit| make_place(hit.doc_type, hit.source))
        .collect())
}

pub fn autocomplete(q: &str,
                    pt_dataset: &Option<&str>,
                    all_data: bool,
                    offset: u64,
                    limit: u64,
                    coord: Option<model::Coord>,
                    cnx: &str,
                    shape: Option<Vec<(f64, f64)>>,
                    types: Option<Vec<&str>>)
                    -> Result<Vec<mimir::Place>, rs_es::error::EsError> {
    fn make_shape(shape: &Option<Vec<(f64, f64)>>) -> Option<Vec<rs_es::units::Location>> {
        shape.as_ref().map(|v| v.iter().map(|&l| l.into()).collect())
    }

    // First we try a prety exact match on the prefix.
    // If there are no results then we do a new fuzzy search (matching ngrams)
    let results = try!(query(&q,
                             &pt_dataset,
                             all_data,
                             cnx,
                             MatchType::Prefix,
                             offset,
                             limit,
                             &coord,
                             make_shape(&shape),
                             &types));
    if results.is_empty() {
        query(&q,
              &pt_dataset,
              all_data,
              cnx,
              MatchType::Fuzzy,
              offset,
              limit,
              &coord,
              make_shape(&shape),
              &types)
    } else {
        Ok(results)
    }
}

#[test]
fn test_make_indexes_impl() {
    // all_data
    assert_eq!(make_indexes_impl(true, &None, &None, |_index| true),
               vec!["munin"]);

    // no dataset and no types
    assert_eq!(make_indexes_impl(false, &None, &None, |_index| true),
               vec!["munin_geo_data"]);

    // dataset fr + no types
    assert_eq!(make_indexes_impl(false,
                                 &Some("munin_stop_fr".to_string()),
                                 &None,
                                 |_index| true),
               vec!["munin_geo_data", "munin_stop_fr"]);

    // no dataset + types poi, city, street, house and public_transport:stop_area
    // => munin_stop is not included
    assert_eq!(make_indexes_impl(false,
                                 &None,
                                 &Some(vec!["poi",
                                            "city",
                                            "street",
                                            "house",
                                            "public_transport:stop_area"]),
                                 |_index| true),
               vec!["munin_poi", "munin_admin", "munin_street", "munin_addr"]);

    // no dataset fr + type public_transport:stop_area only
    assert_eq!(make_indexes_impl(false,
                                 &None,
                                 &Some(vec!["public_transport:stop_area"]),
                                 |_index| true),
               Vec::<String>::new());

    // dataset fr + types poi, city, street, house and public_transport:stop_area
    assert_eq!(make_indexes_impl(false,
                                 &Some("munin_stop_fr".to_string()),
                                 &Some(vec!["poi",
                                            "city",
                                            "street",
                                            "house",
                                            "public_transport:stop_area"]),
                                 |_index| true),
               vec!["munin_poi", "munin_admin", "munin_street", "munin_addr", "munin_stop_fr"]);

    // dataset fr types poi, city, street, house without public_transport:stop_area
    //  => munin_stop_fr is not included
    assert_eq!(make_indexes_impl(false,
                                 &Some("munin_stop_fr".to_string()),
                                 &Some(vec!["poi", "city", "street", "house"]),
                                 |_index| true),
               vec!["munin_poi", "munin_admin", "munin_street", "munin_addr"]);
}
<|endoftext|>


<commit_before>use nom::{Consumer, FileProducer};
use std::io::{Error, ErrorKind, Result};
use std::u32;

use metadata::{
  Block, BlockData,
  StreamInfo, CueSheet, VorbisComment, Picture,
  PictureType,
  MetaDataConsumer,
};

macro_rules! optional_eq (
  ($compare: expr, $option: expr) => (
    $option.map_or(true, |compare| $compare == compare);
  );
);

pub fn get_metadata(filename: &str) -> Result<Vec<Block>> {
  FileProducer::new(filename, 1024).and_then(|mut producer| {
    let mut consumer = MetaDataConsumer::new();

    consumer.run(&mut producer);

    if !consumer.data.is_empty() {
      Ok(consumer.data)
    } else {
      let error_str = "parser: couldn't find any metadata";

      Err(Error::new(ErrorKind::InvalidData, error_str))
    }
  })
}

pub fn get_stream_info(filename: &str) -> Result<StreamInfo> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find StreamInfo";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    for block in blocks {
      if let BlockData::StreamInfo(stream_info) = block.data {
        result = Ok(stream_info);
        break;
      }
    }

    result
  })
}

pub fn get_vorbis_comment(filename: &str) -> Result<VorbisComment> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find VorbisComment";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    for block in blocks {
      if let BlockData::VorbisComment(vorbis_comment) = block.data {
        result = Ok(vorbis_comment);
        break;
      }
    }

    result
  })
}

pub fn get_cue_sheet(filename: &str) -> Result<CueSheet> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find CueSheet";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    for block in blocks {
      if let BlockData::CueSheet(cue_sheet) = block.data {
        result = Ok(cue_sheet);
        break;
      }
    }

    result
  })
}

pub fn get_picture(filename: &str,
                   picture_type: Option<PictureType>,
                   mime_type: Option<&str>,
                   description: Option<&str>,
                   max_width: Option<u32>,
                   max_height: Option<u32>,
                   max_depth: Option<u32>,
                   max_colors: Option<u32>)
                   -> Result<Picture> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find any Picture";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    let mut max_area_seen  = 0;
    let mut max_depth_seen = 0;

    let max_value      = u32::max_value();
    let max_width_num  = max_width.unwrap_or(max_value);
    let max_height_num = max_height.unwrap_or(max_value);
    let max_depth_num  = max_depth.unwrap_or(max_value);
    let max_colors_num = max_colors.unwrap_or(max_value);

    for block in blocks {
      if let BlockData::Picture(picture) = block.data {
        let area = (picture.width as u64) * (picture.height as u64);

        if optional_eq!(picture.picture_type, picture_type) &&
           optional_eq!(picture.mime_type, mime_type) &&
           optional_eq!(picture.description, description) &&
           picture.width <= max_width_num &&
           picture.height <= max_height_num &&
           picture.depth <= max_depth_num &&
           picture.colors <= max_colors_num &&
           (area > max_area_seen || (area == max_area_seen &&
                                     picture.depth > max_depth_seen)) {
          max_area_seen  = area;
          max_depth_seen = picture.depth;
          result         = Ok(picture);
        }
      }
    }

    result
  })
}

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

  #[test]
  #[should_panic]
  fn test_panic_optional_eq() {
    assert!(optional_eq!(0, Some(1)));
  }

  #[test]
  fn test_optional_eq() {
    assert!(optional_eq!(0, None), "Should always return true when None");
    assert!(optional_eq!(0, Some(0)), "Should return true (Some(0) == 0)");
  }

  #[test]
  fn test_get_metadata() {
    let not_found    = get_metadata("non-existent/file.txt");
    let invalid_data = get_metadata("README.md");
    let result       = get_metadata("tests/assets/input-SVAUP.flac");

    assert_eq!(not_found.unwrap_err().kind(), ErrorKind::NotFound);
    assert_eq!(invalid_data.unwrap_err().kind(), ErrorKind::InvalidData);
    assert!(result.is_ok());
  }
}
<commit_msg>Add documentation for public functions<commit_after>use nom::{Consumer, FileProducer};
use std::io::{Error, ErrorKind, Result};
use std::u32;

use metadata::{
  Block, BlockData,
  StreamInfo, CueSheet, VorbisComment, Picture,
  PictureType,
  MetaDataConsumer,
};

macro_rules! optional_eq (
  ($compare: expr, $option: expr) => (
    $option.map_or(true, |compare| $compare == compare);
  );
);

pub fn get_metadata(filename: &str) -> Result<Vec<Block>> {
  FileProducer::new(filename, 1024).and_then(|mut producer| {
    let mut consumer = MetaDataConsumer::new();

    consumer.run(&mut producer);

    if !consumer.data.is_empty() {
      Ok(consumer.data)
    } else {
      let error_str = "parser: couldn't find any metadata";

      Err(Error::new(ErrorKind::InvalidData, error_str))
    }
  })
}

/// Reads and returns the `StreamInfo` metadata block of the given FLAC
/// file.
///
/// # Failures
///
/// * `ErrorKind::NotFound` is returned when the given filename isn't found
///   or there is no `StreamInfo` within the file.
/// * `ErrorKind::InvalidData` is returned when the data within the file
///   isn't valid FLAC data.
pub fn get_stream_info(filename: &str) -> Result<StreamInfo> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find StreamInfo";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    for block in blocks {
      if let BlockData::StreamInfo(stream_info) = block.data {
        result = Ok(stream_info);
        break;
      }
    }

    result
  })
}

/// Reads and returns the `VorbisComment` metadata block of the given FLAC
/// file.
///
/// # Failures
///
/// * `ErrorKind::NotFound` is returned when the given filename isn't found
///   or there is no `VorbisComment` within the file.
/// * `ErrorKind::InvalidData` is returned when the data within the file
///   isn't valid FLAC data.
pub fn get_vorbis_comment(filename: &str) -> Result<VorbisComment> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find VorbisComment";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    for block in blocks {
      if let BlockData::VorbisComment(vorbis_comment) = block.data {
        result = Ok(vorbis_comment);
        break;
      }
    }

    result
  })
}

/// Reads and returns the `CueSheet` metadata block of the given FLAC file.
///
/// # Failures
///
/// * `ErrorKind::NotFound` is returned when the given filename isn't found
///   or there is no `CueSheet` within the file.
/// * `ErrorKind::InvalidData` is returned when the data within the file
///   isn't valid FLAC data.
pub fn get_cue_sheet(filename: &str) -> Result<CueSheet> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find CueSheet";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    for block in blocks {
      if let BlockData::CueSheet(cue_sheet) = block.data {
        result = Ok(cue_sheet);
        break;
      }
    }

    result
  })
}

/// Reads and returns a `Picture` metadata block of the given FLAC file.
///
/// There can be more than one `Picture` block in a file, this function
/// takes optional, that being `Option<T>`, parameters that act as
/// constraints to search within. The `Picture` with the largest area
/// matching all constraints will be returned.
///
/// # Failures
///
/// * `ErrorKind::NotFound` is returned when the given filename isn't found,
///   there is no `Picture` within the file, or no `Picture` that fits the
///   given constraints.
/// * `ErrorKind::InvalidData` is returned when the data within the file
///   isn't valid FLAC data.
pub fn get_picture(filename: &str,
                   picture_type: Option<PictureType>,
                   mime_type: Option<&str>,
                   description: Option<&str>,
                   max_width: Option<u32>,
                   max_height: Option<u32>,
                   max_depth: Option<u32>,
                   max_colors: Option<u32>)
                   -> Result<Picture> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find any Picture";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    let mut max_area_seen  = 0;
    let mut max_depth_seen = 0;

    let max_value      = u32::max_value();
    let max_width_num  = max_width.unwrap_or(max_value);
    let max_height_num = max_height.unwrap_or(max_value);
    let max_depth_num  = max_depth.unwrap_or(max_value);
    let max_colors_num = max_colors.unwrap_or(max_value);

    for block in blocks {
      if let BlockData::Picture(picture) = block.data {
        let area = (picture.width as u64) * (picture.height as u64);

        if optional_eq!(picture.picture_type, picture_type) &&
           optional_eq!(picture.mime_type, mime_type) &&
           optional_eq!(picture.description, description) &&
           picture.width <= max_width_num &&
           picture.height <= max_height_num &&
           picture.depth <= max_depth_num &&
           picture.colors <= max_colors_num &&
           (area > max_area_seen || (area == max_area_seen &&
                                     picture.depth > max_depth_seen)) {
          max_area_seen  = area;
          max_depth_seen = picture.depth;
          result         = Ok(picture);
        }
      }
    }

    result
  })
}

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

  #[test]
  #[should_panic]
  fn test_panic_optional_eq() {
    assert!(optional_eq!(0, Some(1)));
  }

  #[test]
  fn test_optional_eq() {
    assert!(optional_eq!(0, None), "Should always return true when None");
    assert!(optional_eq!(0, Some(0)), "Should return true (Some(0) == 0)");
  }

  #[test]
  fn test_get_metadata() {
    let not_found    = get_metadata("non-existent/file.txt");
    let invalid_data = get_metadata("README.md");
    let result       = get_metadata("tests/assets/input-SVAUP.flac");

    assert_eq!(not_found.unwrap_err().kind(), ErrorKind::NotFound);
    assert_eq!(invalid_data.unwrap_err().kind(), ErrorKind::InvalidData);
    assert!(result.is_ok());
  }
}
<|endoftext|>


<commit_before>#![feature(ascii_ctype)]
use std::ascii::AsciiExt;

pub fn is_pangram(words: &str) -> bool {
    let score = &mut [false; 26];

    for v in words.chars() {
        if !v.is_ascii_alphabetic() {
            continue;
        }

        score[v.to_ascii_lowercase() as usize - 'a' as usize] = true;
    }

    score.iter().all(|&x| x)
}
<commit_msg>Clean up code<commit_after>#![feature(ascii_ctype)]
use std::ascii::AsciiExt;

pub fn is_pangram(words: &str) -> bool {
    let score = &mut [false; 26];

    for v in words.chars() {
        if !v.is_ascii_alphabetic() {
            continue;
        }

        let index = v.to_ascii_lowercase() as usize - 'a' as usize;

        score[index] = true;
    }

    score.iter().all(|&x| x)
}
<|endoftext|>


<commit_before><commit_msg>Fix reading RIP<commit_after><|endoftext|>


<commit_before>#[link(name = "licenseRs", vers = "0.2", author = "Louis Roché")]

extern mod std;
use std::getopts::*;
use std::time::now;
use core::container::Map;
use core::hashmap::linear::LinearMap;
use core::io::ReaderUtil;

fn print_usage(program: &str, _opts: &[std::getopts::Opt]) {
    io::println(fmt!("Usage: %s [options]", program));
    io::println("-h --help\tUsage");
    io::println("--proj\t\tProject name, defaults to name of current directory");
    io::println("--year\t\tCopyright year");
    io::println("--org\t\tOrganization, defaults to $USER");
    io::println("--license\tThe license to generate, one of: agpl3, apache, bsd2, \
                 bsd3, cddl, cc0, epl, gpl2, gpl3, lgpl, mit, mpl");
    io::println("--template\tPath to license template file");
    io::println("--vars\t\tList template variables for specified license");
}

fn get_license(args: Option<~str>) -> ~str {
    match args {
        Some(x) => x.to_lower(),
        None => ~"bsd3"
    }
}

fn get_proj(args: Option<~str>) -> ~str {
    match args {
        Some(x) => x,
        None => {
            match os::getcwd().filename() {
                Some(dir) => dir,
                None => ~"project"
            }
        }
    }
}

fn get_org(args: Option<~str>) -> ~str {
    match args {
        Some(x) => x,
            None => {
                match os::getenv("USER") {
                    Some(dir) => dir,
                    None => ~"organization"
                }
            }
        }
}

fn get_year(args: Option<~str>) -> ~str {
    match args {
        Some(x) => x,
        None => {
            let time = now();
            time.strftime("%Y")
        }
    }
}

fn get_template_path(args: Option<~str>, context: @ mut LinearMap<~str, ~str>) -> ~str {
    match args {
        Some(x) => x,
        None => {
            ~"template-" + *context.get(&~"license") + ~".txt"
        }
    }
}

fn load_file_template(file: ~str) -> ~[u8] {
    let open_file = &io::file_reader(&Path(file));
    if result::is_err(open_file) {
        fail!(~"Cannot open license file");
    }
    let license = result::get(open_file);
    license.read_whole_stream()
}

fn replace_keys(origin_template: ~[u8], context: @ mut LinearMap<~str, ~str>) {
    let mut template = str::from_bytes(origin_template);
    let keys : [~str * 3] = [~"year", ~"project", ~"organization"];
    for keys.each |k| {
        template = str::replace(template, ~"{{ " + *k + ~" }}", *context.get(k));
    }
    io::println(template);
}

fn get_template_vars(template: ~str) -> ~[~str] {
    let vars_tmp = str::split_str_nonempty(template, ~"{{");
    let mut vars: ~[~str] = ~[];
    for vars_tmp.each |&var| {
        let end = str::find_str(var, "}}");
        match end {
            Some(pos) => { vars.push(var.substr(0, pos).trim()) },
            None => { }
        }
    }
    vars
}

fn display_template_vars(origin_template: ~[u8]) {
    let mut template = str::from_bytes(origin_template);
    let vars = get_template_vars(template);
    for vars.each |&var| {
        io::println(var);
    }
}

fn main() {
    let args = os::args();

    let program = copy args[0];

    let opts = ~[
        optflag(~"h"),
        optflag(~"help"),
        optflag(~"vars"),
        optopt(~"year"),
        optopt(~"proj"),
        optopt(~"org"),
        optopt(~"template"),
        optopt(~"license")
    ];
    let matches = match getopts(vec::tail(args), opts) {
        result::Ok(m) => { m }
        result::Err(f) => { fail!(fail_str(f)) }
    };
    if opt_present(&matches, "h") || opt_present(&matches, "help") {
        print_usage(program, opts);
        return;
    }
    let year = opt_maybe_str(&matches, ~"year");
    let proj = opt_maybe_str(&matches, ~"proj");
    let org = opt_maybe_str(&matches, ~"org");
    let license = opt_maybe_str(&matches, ~"license");
    let template_arg = opt_maybe_str(&matches, ~"template");

    let context: @ mut LinearMap<~str, ~str> = @ mut LinearMap::new();
    context.insert(~"year", get_year(year));
    context.insert(~"project", get_proj(proj));
    context.insert(~"organization", get_org(org));
    context.insert(~"license", get_license(license));

    let template_path = get_template_path(template_arg, context);
    let template = load_file_template(template_path);

    if opt_present(&matches, "vars") {
        display_template_vars(template);
    } else {
        replace_keys(template, context);
    }
}<commit_msg>add header option<commit_after>#[link(name = "licenseRs", vers = "0.2", author = "Louis Roché")]

extern mod std;
use std::getopts::*;
use std::time::now;
use core::container::Map;
use core::hashmap::linear::LinearMap;
use core::io::ReaderUtil;

fn print_usage(program: &str, _opts: &[std::getopts::Opt]) {
    io::println(fmt!("Usage: %s [options]", program));
    io::println("-h --help\tUsage");
    io::println("--proj\t\tProject name, defaults to name of current directory");
    io::println("--year\t\tCopyright year");
    io::println("--org\t\tOrganization, defaults to $USER");
    io::println("--license\tThe license to generate, one of: agpl3, apache, bsd2, \
                 bsd3, cddl, cc0, epl, gpl2, gpl3, lgpl, mit, mpl");
    io::println("--template\tPath to license template file");
    io::println("--header\tGenerate header file instead of the license");
    io::println("--vars\t\tList template variables for specified license");
}

fn get_license(args: Option<~str>) -> ~str {
    match args {
        Some(x) => x.to_lower(),
        None => ~"bsd3"
    }
}

fn get_proj(args: Option<~str>) -> ~str {
    match args {
        Some(x) => x,
        None => {
            match os::getcwd().filename() {
                Some(dir) => dir,
                None => ~"project"
            }
        }
    }
}

fn get_org(args: Option<~str>) -> ~str {
    match args {
        Some(x) => x,
            None => {
                match os::getenv("USER") {
                    Some(dir) => dir,
                    None => ~"organization"
                }
            }
        }
}

fn get_year(args: Option<~str>) -> ~str {
    match args {
        Some(x) => x,
        None => {
            let time = now();
            time.strftime("%Y")
        }
    }
}

fn get_template_path(args: Option<~str>, context: @ mut LinearMap<~str, ~str>, header: bool) -> ~str {
    match args {
        Some(x) => x,
        None => {
            let license = context.get(&~"license");
            if header {
                ~"template-" + *license + ~"-header.txt"
            } else {
                ~"template-" + *license + ~".txt"
            }
        }
    }
}

fn load_file_template(file: ~str) -> ~[u8] {
    let open_file = &io::file_reader(&Path(file));
    if result::is_err(open_file) {
        fail!(~"Cannot open license file");
    }
    let license = result::get(open_file);
    license.read_whole_stream()
}

fn replace_keys(origin_template: ~[u8], context: @ mut LinearMap<~str, ~str>) {
    let mut template = str::from_bytes(origin_template);
    let keys : [~str * 3] = [~"year", ~"project", ~"organization"];
    for keys.each |k| {
        template = str::replace(template, ~"{{ " + *k + ~" }}", *context.get(k));
    }
    io::println(template);
}

fn get_template_vars(template: ~str) -> ~[~str] {
    let vars_tmp = str::split_str_nonempty(template, ~"{{");
    let mut vars: ~[~str] = ~[];
    for vars_tmp.each |&var| {
        let end = str::find_str(var, "}}");
        match end {
            Some(pos) => {
                let var_name = var.substr(0, pos).trim();
                if !vars.contains(&var_name) {
                    vars.push(var_name);
                }
            },
            None => { }
        }
    }
    vars
}

fn display_template_vars(origin_template: ~[u8]) {
    let mut template = str::from_bytes(origin_template);
    let vars = get_template_vars(template);
    for vars.each |&var| {
        io::println(var);
    }
}

fn main() {
    let args = os::args();

    let program = copy args[0];

    let opts = ~[
        optflag(~"h"),
        optflag(~"help"),
        optflag(~"vars"),
        optflag(~"header"),
        optopt(~"year"),
        optopt(~"proj"),
        optopt(~"org"),
        optopt(~"template"),
        optopt(~"license")
    ];
    let matches = match getopts(vec::tail(args), opts) {
        result::Ok(m) => { m }
        result::Err(f) => { fail!(fail_str(f)) }
    };
    if opt_present(&matches, "h") || opt_present(&matches, "help") {
        print_usage(program, opts);
        return;
    }
    let year = opt_maybe_str(&matches, ~"year");
    let proj = opt_maybe_str(&matches, ~"proj");
    let org = opt_maybe_str(&matches, ~"org");
    let license = opt_maybe_str(&matches, ~"license");
    let template_arg = opt_maybe_str(&matches, ~"template");
    let header = opt_present(&matches, "header");

    let context: @ mut LinearMap<~str, ~str> = @ mut LinearMap::new();
    context.insert(~"year", get_year(year));
    context.insert(~"project", get_proj(proj));
    context.insert(~"organization", get_org(org));
    context.insert(~"license", get_license(license));

    let template_path = get_template_path(template_arg, context, header);
    let template = load_file_template(template_path);

    if opt_present(&matches, "vars") {
        display_template_vars(template);
    } else {
        replace_keys(template, context);
    }
}<|endoftext|>


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

//! An implementation of SipHash 2-4.

use prelude::v1::*;

use ptr;
use super::Hasher;

/// An implementation of SipHash 2-4.
///
/// See: http://131002.net/siphash/
///
/// Consider this as a main "general-purpose" hash for all hashtables: it
/// runs at good speed (competitive with spooky and city) and permits
/// strong _keyed_ hashing. Key your hashtables from a strong RNG,
/// such as `rand::Rng`.
///
/// Although the SipHash algorithm is considered to be cryptographically
/// strong, this implementation has not been reviewed for such purposes.
/// As such, all cryptographic uses of this implementation are strongly
/// discouraged.
#[stable(feature = "rust1", since = "1.0.0")]
pub struct SipHasher {
    k0: u64,
    k1: u64,
    length: usize, // how many bytes we've processed
    // v0, v2 and v1, v3 show up in pairs in the algorithm,
    // and simd implementations of SipHash will use vectors
    // of v02 and v13. By placing them in this order in the struct,
    // the compiler can pick up on just a few simd optimizations by itself.
    v0: u64, // hash state
    v2: u64,
    v1: u64,
    v3: u64,
    tail: u64, // unprocessed bytes le
    ntail: usize, // how many bytes in tail are valid
}

// sadly, these macro definitions can't appear later,
// because they're needed in the following defs;
// this design could be improved.

macro_rules! u8to64_le {
    ($buf:expr, $i:expr) =>
    ($buf[0+$i] as u64 |
     ($buf[1+$i] as u64) << 8 |
     ($buf[2+$i] as u64) << 16 |
     ($buf[3+$i] as u64) << 24 |
     ($buf[4+$i] as u64) << 32 |
     ($buf[5+$i] as u64) << 40 |
     ($buf[6+$i] as u64) << 48 |
     ($buf[7+$i] as u64) << 56);
    ($buf:expr, $i:expr, $len:expr) =>
    ({
        let mut t = 0;
        let mut out = 0;
        while t < $len {
            out |= ($buf[t+$i] as u64) << t*8;
            t += 1;
        }
        out
    });
}

/// Load a full u64 word from a byte stream, in LE order. Use
/// `copy_nonoverlapping` to let the compiler generate the most efficient way
/// to load u64 from a possibly unaligned address.
///
/// Unsafe because: unchecked indexing at i..i+8
#[inline]
unsafe fn load_u64_le(buf: &[u8], i: usize) -> u64 {
    debug_assert!(i + 8 <= buf.len());
    let mut data = 0u64;
    ptr::copy_nonoverlapping(buf.get_unchecked(i), &mut data as *mut _ as *mut u8, 8);
    data.to_le()
}

macro_rules! rotl {
    ($x:expr, $b:expr) =>
    (($x << $b) | ($x >> (64_i32.wrapping_sub($b))))
}

macro_rules! compress {
    ($v0:expr, $v1:expr, $v2:expr, $v3:expr) =>
    ({
        $v0 = $v0.wrapping_add($v1); $v1 = rotl!($v1, 13); $v1 ^= $v0;
        $v0 = rotl!($v0, 32);
        $v2 = $v2.wrapping_add($v3); $v3 = rotl!($v3, 16); $v3 ^= $v2;
        $v0 = $v0.wrapping_add($v3); $v3 = rotl!($v3, 21); $v3 ^= $v0;
        $v2 = $v2.wrapping_add($v1); $v1 = rotl!($v1, 17); $v1 ^= $v2;
        $v2 = rotl!($v2, 32);
    })
}

impl SipHasher {
    /// Creates a new `SipHasher` with the two initial keys set to 0.
    #[inline]
    #[stable(feature = "rust1", since = "1.0.0")]
    pub fn new() -> SipHasher {
        SipHasher::new_with_keys(0, 0)
    }

    /// Creates a `SipHasher` that is keyed off the provided keys.
    #[inline]
    #[stable(feature = "rust1", since = "1.0.0")]
    pub fn new_with_keys(key0: u64, key1: u64) -> SipHasher {
        let mut state = SipHasher {
            k0: key0,
            k1: key1,
            length: 0,
            v0: 0,
            v1: 0,
            v2: 0,
            v3: 0,
            tail: 0,
            ntail: 0,
        };
        state.reset();
        state
    }

    #[inline]
    fn reset(&mut self) {
        self.length = 0;
        self.v0 = self.k0 ^ 0x736f6d6570736575;
        self.v1 = self.k1 ^ 0x646f72616e646f6d;
        self.v2 = self.k0 ^ 0x6c7967656e657261;
        self.v3 = self.k1 ^ 0x7465646279746573;
        self.ntail = 0;
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl Hasher for SipHasher {
    #[inline]
    fn write(&mut self, msg: &[u8]) {
        let length = msg.len();
        self.length += length;

        let mut needed = 0;

        if self.ntail != 0 {
            needed = 8 - self.ntail;
            if length < needed {
                self.tail |= u8to64_le!(msg, 0, length) << 8 * self.ntail;
                self.ntail += length;
                return
            }

            let m = self.tail | u8to64_le!(msg, 0, needed) << 8 * self.ntail;

            self.v3 ^= m;
            compress!(self.v0, self.v1, self.v2, self.v3);
            compress!(self.v0, self.v1, self.v2, self.v3);
            self.v0 ^= m;

            self.ntail = 0;
        }

        // Buffered tail is now flushed, process new input.
        let len = length - needed;
        let left = len & 0x7;

        let mut i = needed;
        while i < len - left {
            let mi = unsafe { load_u64_le(msg, i) };

            self.v3 ^= mi;
            compress!(self.v0, self.v1, self.v2, self.v3);
            compress!(self.v0, self.v1, self.v2, self.v3);
            self.v0 ^= mi;

            i += 8;
        }

        self.tail = u8to64_le!(msg, i, left);
        self.ntail = left;
    }

    #[inline]
    fn finish(&self) -> u64 {
        let mut v0 = self.v0;
        let mut v1 = self.v1;
        let mut v2 = self.v2;
        let mut v3 = self.v3;

        let b: u64 = ((self.length as u64 & 0xff) << 56) | self.tail;

        v3 ^= b;
        compress!(v0, v1, v2, v3);
        compress!(v0, v1, v2, v3);
        v0 ^= b;

        v2 ^= 0xff;
        compress!(v0, v1, v2, v3);
        compress!(v0, v1, v2, v3);
        compress!(v0, v1, v2, v3);
        compress!(v0, v1, v2, v3);

        v0 ^ v1 ^ v2 ^ v3
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl Clone for SipHasher {
    #[inline]
    fn clone(&self) -> SipHasher {
        SipHasher {
            k0: self.k0,
            k1: self.k1,
            length: self.length,
            v0: self.v0,
            v1: self.v1,
            v2: self.v2,
            v3: self.v3,
            tail: self.tail,
            ntail: self.ntail,
        }
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl Default for SipHasher {
    fn default() -> SipHasher {
        SipHasher::new()
    }
}
<commit_msg>Rollup merge of #32091 - dirk:dirk/siphasher-docs-clarification, r=alexcrichton<commit_after>// Copyright 2012-2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! An implementation of SipHash 2-4.

use prelude::v1::*;

use ptr;
use super::Hasher;

/// An implementation of SipHash 2-4.
///
/// See: http://131002.net/siphash/
///
/// This is currently the default hashing function used by standard library
/// (eg. `collections::HashMap` uses it by default).
///
/// SipHash is a general-purpose hashing function: it runs at a good
/// speed (competitive with Spooky and City) and permits strong _keyed_
/// hashing. This lets you key your hashtables from a strong RNG, such as
/// [`rand::os::OsRng`](https://doc.rust-lang.org/rand/rand/os/struct.OsRng.html).
///
/// Although the SipHash algorithm is considered to be generally strong,
/// it is not intended for cryptographic purposes. As such, all
/// cryptographic uses of this implementation are _strongly discouraged_.
#[stable(feature = "rust1", since = "1.0.0")]
pub struct SipHasher {
    k0: u64,
    k1: u64,
    length: usize, // how many bytes we've processed
    // v0, v2 and v1, v3 show up in pairs in the algorithm,
    // and simd implementations of SipHash will use vectors
    // of v02 and v13. By placing them in this order in the struct,
    // the compiler can pick up on just a few simd optimizations by itself.
    v0: u64, // hash state
    v2: u64,
    v1: u64,
    v3: u64,
    tail: u64, // unprocessed bytes le
    ntail: usize, // how many bytes in tail are valid
}

// sadly, these macro definitions can't appear later,
// because they're needed in the following defs;
// this design could be improved.

macro_rules! u8to64_le {
    ($buf:expr, $i:expr) =>
    ($buf[0+$i] as u64 |
     ($buf[1+$i] as u64) << 8 |
     ($buf[2+$i] as u64) << 16 |
     ($buf[3+$i] as u64) << 24 |
     ($buf[4+$i] as u64) << 32 |
     ($buf[5+$i] as u64) << 40 |
     ($buf[6+$i] as u64) << 48 |
     ($buf[7+$i] as u64) << 56);
    ($buf:expr, $i:expr, $len:expr) =>
    ({
        let mut t = 0;
        let mut out = 0;
        while t < $len {
            out |= ($buf[t+$i] as u64) << t*8;
            t += 1;
        }
        out
    });
}

/// Load a full u64 word from a byte stream, in LE order. Use
/// `copy_nonoverlapping` to let the compiler generate the most efficient way
/// to load u64 from a possibly unaligned address.
///
/// Unsafe because: unchecked indexing at i..i+8
#[inline]
unsafe fn load_u64_le(buf: &[u8], i: usize) -> u64 {
    debug_assert!(i + 8 <= buf.len());
    let mut data = 0u64;
    ptr::copy_nonoverlapping(buf.get_unchecked(i), &mut data as *mut _ as *mut u8, 8);
    data.to_le()
}

macro_rules! rotl {
    ($x:expr, $b:expr) =>
    (($x << $b) | ($x >> (64_i32.wrapping_sub($b))))
}

macro_rules! compress {
    ($v0:expr, $v1:expr, $v2:expr, $v3:expr) =>
    ({
        $v0 = $v0.wrapping_add($v1); $v1 = rotl!($v1, 13); $v1 ^= $v0;
        $v0 = rotl!($v0, 32);
        $v2 = $v2.wrapping_add($v3); $v3 = rotl!($v3, 16); $v3 ^= $v2;
        $v0 = $v0.wrapping_add($v3); $v3 = rotl!($v3, 21); $v3 ^= $v0;
        $v2 = $v2.wrapping_add($v1); $v1 = rotl!($v1, 17); $v1 ^= $v2;
        $v2 = rotl!($v2, 32);
    })
}

impl SipHasher {
    /// Creates a new `SipHasher` with the two initial keys set to 0.
    #[inline]
    #[stable(feature = "rust1", since = "1.0.0")]
    pub fn new() -> SipHasher {
        SipHasher::new_with_keys(0, 0)
    }

    /// Creates a `SipHasher` that is keyed off the provided keys.
    #[inline]
    #[stable(feature = "rust1", since = "1.0.0")]
    pub fn new_with_keys(key0: u64, key1: u64) -> SipHasher {
        let mut state = SipHasher {
            k0: key0,
            k1: key1,
            length: 0,
            v0: 0,
            v1: 0,
            v2: 0,
            v3: 0,
            tail: 0,
            ntail: 0,
        };
        state.reset();
        state
    }

    #[inline]
    fn reset(&mut self) {
        self.length = 0;
        self.v0 = self.k0 ^ 0x736f6d6570736575;
        self.v1 = self.k1 ^ 0x646f72616e646f6d;
        self.v2 = self.k0 ^ 0x6c7967656e657261;
        self.v3 = self.k1 ^ 0x7465646279746573;
        self.ntail = 0;
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl Hasher for SipHasher {
    #[inline]
    fn write(&mut self, msg: &[u8]) {
        let length = msg.len();
        self.length += length;

        let mut needed = 0;

        if self.ntail != 0 {
            needed = 8 - self.ntail;
            if length < needed {
                self.tail |= u8to64_le!(msg, 0, length) << 8 * self.ntail;
                self.ntail += length;
                return
            }

            let m = self.tail | u8to64_le!(msg, 0, needed) << 8 * self.ntail;

            self.v3 ^= m;
            compress!(self.v0, self.v1, self.v2, self.v3);
            compress!(self.v0, self.v1, self.v2, self.v3);
            self.v0 ^= m;

            self.ntail = 0;
        }

        // Buffered tail is now flushed, process new input.
        let len = length - needed;
        let left = len & 0x7;

        let mut i = needed;
        while i < len - left {
            let mi = unsafe { load_u64_le(msg, i) };

            self.v3 ^= mi;
            compress!(self.v0, self.v1, self.v2, self.v3);
            compress!(self.v0, self.v1, self.v2, self.v3);
            self.v0 ^= mi;

            i += 8;
        }

        self.tail = u8to64_le!(msg, i, left);
        self.ntail = left;
    }

    #[inline]
    fn finish(&self) -> u64 {
        let mut v0 = self.v0;
        let mut v1 = self.v1;
        let mut v2 = self.v2;
        let mut v3 = self.v3;

        let b: u64 = ((self.length as u64 & 0xff) << 56) | self.tail;

        v3 ^= b;
        compress!(v0, v1, v2, v3);
        compress!(v0, v1, v2, v3);
        v0 ^= b;

        v2 ^= 0xff;
        compress!(v0, v1, v2, v3);
        compress!(v0, v1, v2, v3);
        compress!(v0, v1, v2, v3);
        compress!(v0, v1, v2, v3);

        v0 ^ v1 ^ v2 ^ v3
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl Clone for SipHasher {
    #[inline]
    fn clone(&self) -> SipHasher {
        SipHasher {
            k0: self.k0,
            k1: self.k1,
            length: self.length,
            v0: self.v0,
            v1: self.v1,
            v2: self.v2,
            v3: self.v3,
            tail: self.tail,
            ntail: self.ntail,
        }
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl Default for SipHasher {
    fn default() -> SipHasher {
        SipHasher::new()
    }
}
<|endoftext|>