<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

#[macro_use]
extern crate log;
extern crate osmpbfreader;
extern crate rustc_serialize;
extern crate docopt;
extern crate mimir;
extern crate rs_es;

use std::collections::HashSet;
use std::collections::HashMap;
use mimirsbrunn::rubber::Rubber;
use mimirsbrunn::objects::{Polygon, MultiPolygon};

pub type AdminsVec = Vec<mimir::Admin>;
pub type StreetsVec = Vec<mimir::Street>;
pub type ParsedPbf = osmpbfreader::OsmPbfReader<std::fs::File>;

#[derive(RustcDecodable, Debug)]
struct Args {
    flag_input: String,
    flag_level: Vec<u32>,
    flag_connection_string: String,
}

static USAGE: &'static str = "
Usage:
    osm2mimir --help
    osm2mimir --input=<file> [--connection-string=<connection-string>] --level=<level>...

Options:
    -h, --help            Show this message.
    -i, --input=<file>    OSM PBF file.
    -l, --level=<level>   Admin levels to keep.
    -c, --connection-string=<connection-string>
                          Elasticsearch parameters, [default: http://localhost:9200/munin]
";

#[derive(Debug)]
struct AdminMatcher {
    admin_levels: HashSet<u32>,
}
impl AdminMatcher {
    pub fn new(levels: HashSet<u32>) -> AdminMatcher {
        AdminMatcher { admin_levels: levels }
    }

    pub fn is_admin(&self, obj: &osmpbfreader::OsmObj) -> bool {
        match *obj {
            osmpbfreader::OsmObj::Relation(ref rel) => {
                rel.tags.get("boundary").map_or(false, |v| v == "administrative") &&
                rel.tags.get("admin_level").map_or(false, |lvl| {
                    self.admin_levels.contains(&lvl.parse::<u32>().unwrap_or(0))
                })
            }
            _ => false,
        }
    }
}

struct BoundaryPart {
    nodes: Vec<osmpbfreader::Node>,
}

impl BoundaryPart {
    pub fn new(nodes: Vec<osmpbfreader::Node>) -> BoundaryPart {
        BoundaryPart { nodes: nodes }
    }
    pub fn first(&self) -> i64 {
        self.nodes.first().unwrap().id
    }
    pub fn last(&self) -> i64 {
        self.nodes.last().unwrap().id
    }
}

fn get_nodes(way: &osmpbfreader::Way,
             objects: &HashMap<osmpbfreader::OsmId, osmpbfreader::OsmObj>)
             -> Vec<osmpbfreader::Node> {
    way.nodes
       .iter()
       .filter_map(|node_id| objects.get(&osmpbfreader::OsmId::Node(*node_id)))
       .filter_map(|node_obj| if let &osmpbfreader::OsmObj::Node(ref node) = node_obj {
           Some(node.clone())
       } else {
           None
       })
       .collect()
}

fn build_boundary(relation: &osmpbfreader::Relation,
                  objects: &HashMap<osmpbfreader::OsmId, osmpbfreader::OsmObj>)
                  -> Option<MultiPolygon> {
    let mut boundary_parts: Vec<BoundaryPart> =
        relation.refs
                .iter()
                .filter(|rf| rf.role == "outer" || rf.role == "" || rf.role == "enclave")
                .filter_map(|refe| {
                    objects.get(&refe.member).or_else(|| {
                        warn!("missing element for relation {}", relation.id);
                        None
                    })
                })
                .filter_map(|way_obj| if let &osmpbfreader::OsmObj::Way(ref way) = way_obj {
                    Some(way)
                } else {
                    None
                })
                .map(|way| get_nodes(&way, objects))
                .filter(|nodes| nodes.len() > 1)
                .map(|nodes| BoundaryPart::new(nodes))
                .collect();
    let mut multipoly = MultiPolygon { polygons: Vec::new() };
    // we want to try build a polygon for a least each way
    while !boundary_parts.is_empty() {
        let mut current = -1;
        let mut first = 0;
        let mut nb_try = 0;

        let mut outer: Vec<osmpbfreader::Node> = Vec::new();
        // we try to close the polygon, if we can't we want to at least have tried one time per
        // way. We could improve that latter by trying to attach the way to both side of the
        // polygon
        let max_try = boundary_parts.len();
        while current != first && nb_try < max_try {
            let mut i = 0;
            while i < boundary_parts.len() {
                if outer.is_empty() {
                    // our polygon is empty, we initialise it with the current way
                    first = boundary_parts[i].first();
                    current = boundary_parts[i].last();
                    outer.append(&mut boundary_parts[i].nodes);
                    // this way has been used, we remove it from the pool
                    boundary_parts.remove(i);
                    continue;
                }
                if current == boundary_parts[i].first() {
                    // the start of current way touch the polygon, we add it and remove it from the
                    // pool
                    current = boundary_parts[i].last();
                    outer.append(&mut boundary_parts[i].nodes);
                    boundary_parts.remove(i);
                } else if current == boundary_parts[i].last() {
                    // the end of the current way touch the polygon, we reverse the way and add it
                    current = boundary_parts[i].first();
                    boundary_parts[i].nodes.reverse();
                    outer.append(&mut boundary_parts[i].nodes);
                    boundary_parts.remove(i);
                } else {
                    i += 1;
                    // didnt do anything, we want to explore the next way, if we had do something we
                    // will have removed the current way and there will be no need to increment
                }
                if current == first {
                    // our polygon is closed, we create it and add it to the multipolygon
                    let polygon = Polygon::new(outer.iter()
                                                    .map(|n| {
                                                        mimir::Coord {
                                                            lat: n.lat,
                                                            lon: n.lon,
                                                        }
                                                    })
                                                    .collect());
                    multipoly.polygons.push(polygon);
                    break;
                }
            }
            nb_try += 1;
        }
    }
    debug!("polygon for relation {};{}", relation.id, multipoly.to_wkt());
    if multipoly.polygons.is_empty() {
        None
    } else {
        Some(multipoly)
    }
}

fn parse_osm_pbf(path: &String) -> ParsedPbf {
    let path = std::path::Path::new(&path);
    osmpbfreader::OsmPbfReader::new(std::fs::File::open(&path).unwrap())
}

fn administrative_regions(pbf: &mut ParsedPbf, levels: HashSet<u32>) -> AdminsVec {
    let mut administrative_regions = AdminsVec::new();
    let matcher = AdminMatcher::new(levels);
    let objects = osmpbfreader::get_objs_and_deps(pbf, |o| matcher.is_admin(o)).unwrap();
    // load administratives regions
    for (_, obj) in &objects {
        if !matcher.is_admin(&obj) {
            continue;
        }
        if let &osmpbfreader::OsmObj::Relation(ref relation) = obj {
            let level = relation.tags
                                .get("admin_level")
                                .and_then(|s| s.parse().ok());
            let level = match level {
                None => {
                    info!("invalid admin_level for relation {}: admin_level {:?}",
                          relation.id,
                          relation.tags.get("admin_level"));
                    continue;
                }
                Some(l) => l,
            };
            // administrative region with name ?
            let name = match relation.tags.get("name") {
                Some(val) => val,
                None => {
                    warn!("adminstrative region without name for relation {}:  admin_level {} \
                           ignored.",
                          relation.id,
                          level);
                    continue;
                }
            };
            // admininstrative region without coordinates
            let coord_centre = relation.refs
                                       .iter()
                                       .find(|rf| rf.role == "admin_centre")
                                       .and_then(|r| {
                                           objects.get(&r.member).and_then(|value| {
                                               match value {
                                                   &osmpbfreader::OsmObj::Node(ref node) => {
                                                       Some(mimir::Coord {
                                                           lat: node.lat,
                                                           lon: node.lon,
                                                       })
                                                   }
                                                   _ => None,
                                               }
                                           })
                                       });

            let admin_id = match relation.tags.get("ref:INSEE") {
                Some(val) => format!("admin:fr:{}", val.trim_left_matches('0')),
                None => format!("admin:osm:{}", relation.id),
            };
            let zip_code = match relation.tags.get("addr:postcode") {
                Some(val) => &val[..],
                None => "",
            };
            let boundary = build_boundary(&relation, &objects);
            let admin = mimir::Admin {
                id: admin_id,
                level: level,
                name: name.to_string(),
                zip_code: zip_code.to_string(),
                // TODO weight value ?
                weight: 1,
                coord: coord_centre,
                boundary: boundary,
            };
            administrative_regions.push(admin);
        }
    }
    return administrative_regions;
}

fn streets(pbf: &mut ParsedPbf) -> StreetsVec {

    let is_valid_street = |obj: &osmpbfreader::OsmObj| -> bool {
        match *obj {
            osmpbfreader::OsmObj::Way(ref way) => {
                way.tags.get("highway").map_or(false, |x: &String| !x.is_empty()) &&
                way.tags.get("name").map_or(false, |x: &String| !x.is_empty())
            }
            _ => false,
        }
    };

    let is_associated_street_relation = |obj: &osmpbfreader::OsmObj| -> bool {
        match *obj {
            osmpbfreader::OsmObj::Relation(ref rel) => {
                rel.tags.get("type").map_or(false, |v| v == "associatedStreet")
            }
            _ => false,
        }
    };

    let mut street_map = osmpbfreader::get_objs_and_deps(pbf, is_valid_street).unwrap();
    let relation_map = osmpbfreader::get_objs_and_deps(pbf, is_associated_street_relation).unwrap();

    // Sometimes, streets can be devided into several "way"s that still have the same street name.
    // The reason why a street is devided may be that a part of the street become a bridge/tunne/etc.
    // In this case, a "relation" tagged with (type = associatedStreet) is used to group all these "way"s.
    // In order not to have doublons in autocompleion, we should keep only one "way" in the relation
    // and remove all the rest way whose street name is the same.
    for (_, rel_obj) in &relation_map {
        if let &osmpbfreader::OsmObj::Relation(ref rel) = rel_obj {
            let mut found = false;
            for ref_obj in &rel.refs {
                if let &osmpbfreader::OsmId::Way(_) = &ref_obj.member {
                    if !found {
                        found = true;
                        continue;
                    };
                    street_map.remove(&ref_obj.member);
                };
            }
        };
    }

    street_map.iter()
              .filter_map(|(osm_id, obj)| {
                  if let &osmpbfreader::OsmObj::Way(ref way) = obj {
                      if let &osmpbfreader::OsmId::Way(ref way_id) = osm_id {
                          way.tags.get("name").and_then(|way_name| {
                              Some(mimirsbrunn::Street {
                                  id: way_id.to_string(),
                                  street_name: way_name.to_string(),
                                  name: way_name.to_string(),
                                  weight: 1,
                                  administrative_region: None,
                              })
                          })
                      } else {
                          None
                      }
                  } else {
                      None
                  }
              })
              .collect()

}

fn index_osm(es_cnx_string: &str, admins: &AdminsVec, streets: &StreetsVec) {
    let mut rubber = Rubber::new(es_cnx_string);
    rubber.create_index();
    match rubber.clean_db_by_doc_type(&["admin", "street"]) {
        Err(e) => panic!("failed to clean data by document type: {}", e),
        Ok(nb) => info!("clean data by document type : {}", nb),
    }
    info!("Add data in elasticsearch db.");
    match rubber.bulk_index(admins.iter()) {
        Err(e) => panic!("failed to index admins of osm because: {}", e),
        Ok(nb) => info!("Nb of indexed adminstrative regions: {}", nb),
    }

    match rubber.bulk_index(streets.iter()) {
        Err(e) => panic!("failed to index streets of osm because: {}", e),
        Ok(nb) => info!("Nb of indexed streets: {}", nb),
    }
}

fn main() {
    mimir::logger_init().unwrap();
    debug!("importing adminstrative region into Mimir");
    let args: Args = docopt::Docopt::new(USAGE)
                         .and_then(|d| d.decode())
                         .unwrap_or_else(|e| e.exit());

    let levels = args.flag_level.iter().cloned().collect();
    let mut parsed_pbf = parse_osm_pbf(&args.flag_input);
    let admins = administrative_regions(&mut parsed_pbf, levels);
    let streets = streets(&mut parsed_pbf);
    index_osm(&args.flag_connection_string, &admins, &streets);
}
<commit_msg>minor fix<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

#[macro_use]
extern crate log;
extern crate osmpbfreader;
extern crate rustc_serialize;
extern crate docopt;
extern crate mimir;
extern crate rs_es;

use std::collections::HashSet;
use std::collections::HashMap;
use mimirsbrunn::rubber::Rubber;
use mimirsbrunn::objects::{Polygon, MultiPolygon};

pub type AdminsVec = Vec<mimir::Admin>;
pub type StreetsVec = Vec<mimir::Street>;
pub type ParsedPbf = osmpbfreader::OsmPbfReader<std::fs::File>;

#[derive(RustcDecodable, Debug)]
struct Args {
    flag_input: String,
    flag_level: Vec<u32>,
    flag_connection_string: String,
}

static USAGE: &'static str = "
Usage:
    osm2mimir --help
    osm2mimir --input=<file> [--connection-string=<connection-string>] --level=<level>...

Options:
    -h, --help            Show this message.
    -i, --input=<file>    OSM PBF file.
    -l, --level=<level>   Admin levels to keep.
    -c, --connection-string=<connection-string>
                          Elasticsearch parameters, [default: http://localhost:9200/munin]
";

#[derive(Debug)]
struct AdminMatcher {
    admin_levels: HashSet<u32>,
}
impl AdminMatcher {
    pub fn new(levels: HashSet<u32>) -> AdminMatcher {
        AdminMatcher { admin_levels: levels }
    }

    pub fn is_admin(&self, obj: &osmpbfreader::OsmObj) -> bool {
        match *obj {
            osmpbfreader::OsmObj::Relation(ref rel) => {
                rel.tags.get("boundary").map_or(false, |v| v == "administrative") &&
                rel.tags.get("admin_level").map_or(false, |lvl| {
                    self.admin_levels.contains(&lvl.parse::<u32>().unwrap_or(0))
                })
            }
            _ => false,
        }
    }
}

struct BoundaryPart {
    nodes: Vec<osmpbfreader::Node>,
}

impl BoundaryPart {
    pub fn new(nodes: Vec<osmpbfreader::Node>) -> BoundaryPart {
        BoundaryPart { nodes: nodes }
    }
    pub fn first(&self) -> i64 {
        self.nodes.first().unwrap().id
    }
    pub fn last(&self) -> i64 {
        self.nodes.last().unwrap().id
    }
}

fn get_nodes(way: &osmpbfreader::Way,
             objects: &HashMap<osmpbfreader::OsmId, osmpbfreader::OsmObj>)
             -> Vec<osmpbfreader::Node> {
    way.nodes
       .iter()
       .filter_map(|node_id| objects.get(&osmpbfreader::OsmId::Node(*node_id)))
       .filter_map(|node_obj| if let &osmpbfreader::OsmObj::Node(ref node) = node_obj {
           Some(node.clone())
       } else {
           None
       })
       .collect()
}

fn build_boundary(relation: &osmpbfreader::Relation,
                  objects: &HashMap<osmpbfreader::OsmId, osmpbfreader::OsmObj>)
                  -> Option<MultiPolygon> {
    let mut boundary_parts: Vec<BoundaryPart> =
        relation.refs
                .iter()
                .filter(|rf| rf.role == "outer" || rf.role == "" || rf.role == "enclave")
                .filter_map(|refe| {
                    objects.get(&refe.member).or_else(|| {
                        warn!("missing element for relation {}", relation.id);
                        None
                    })
                })
                .filter_map(|way_obj| if let &osmpbfreader::OsmObj::Way(ref way) = way_obj {
                    Some(way)
                } else {
                    None
                })
                .map(|way| get_nodes(&way, objects))
                .filter(|nodes| nodes.len() > 1)
                .map(|nodes| BoundaryPart::new(nodes))
                .collect();
    let mut multipoly = MultiPolygon { polygons: Vec::new() };
    // we want to try build a polygon for a least each way
    while !boundary_parts.is_empty() {
        let mut current = -1;
        let mut first = 0;
        let mut nb_try = 0;

        let mut outer: Vec<osmpbfreader::Node> = Vec::new();
        // we try to close the polygon, if we can't we want to at least have tried one time per
        // way. We could improve that latter by trying to attach the way to both side of the
        // polygon
        let max_try = boundary_parts.len();
        while current != first && nb_try < max_try {
            let mut i = 0;
            while i < boundary_parts.len() {
                if outer.is_empty() {
                    // our polygon is empty, we initialise it with the current way
                    first = boundary_parts[i].first();
                    current = boundary_parts[i].last();
                    outer.append(&mut boundary_parts[i].nodes);
                    // this way has been used, we remove it from the pool
                    boundary_parts.remove(i);
                    continue;
                }
                if current == boundary_parts[i].first() {
                    // the start of current way touch the polygon, we add it and remove it from the
                    // pool
                    current = boundary_parts[i].last();
                    outer.append(&mut boundary_parts[i].nodes);
                    boundary_parts.remove(i);
                } else if current == boundary_parts[i].last() {
                    // the end of the current way touch the polygon, we reverse the way and add it
                    current = boundary_parts[i].first();
                    boundary_parts[i].nodes.reverse();
                    outer.append(&mut boundary_parts[i].nodes);
                    boundary_parts.remove(i);
                } else {
                    i += 1;
                    // didnt do anything, we want to explore the next way, if we had do something we
                    // will have removed the current way and there will be no need to increment
                }
                if current == first {
                    // our polygon is closed, we create it and add it to the multipolygon
                    let polygon = Polygon::new(outer.iter()
                                                    .map(|n| {
                                                        mimir::Coord {
                                                            lat: n.lat,
                                                            lon: n.lon,
                                                        }
                                                    })
                                                    .collect());
                    multipoly.polygons.push(polygon);
                    break;
                }
            }
            nb_try += 1;
        }
    }
    debug!("polygon for relation {};{}", relation.id, multipoly.to_wkt());
    if multipoly.polygons.is_empty() {
        None
    } else {
        Some(multipoly)
    }
}

fn parse_osm_pbf(path: &String) -> ParsedPbf {
    let path = std::path::Path::new(&path);
    osmpbfreader::OsmPbfReader::new(std::fs::File::open(&path).unwrap())
}

fn administrative_regions(pbf: &mut ParsedPbf, levels: HashSet<u32>) -> AdminsVec {
    let mut administrative_regions = AdminsVec::new();
    let matcher = AdminMatcher::new(levels);
    let objects = osmpbfreader::get_objs_and_deps(pbf, |o| matcher.is_admin(o)).unwrap();
    // load administratives regions
    for (_, obj) in &objects {
        if !matcher.is_admin(&obj) {
            continue;
        }
        if let &osmpbfreader::OsmObj::Relation(ref relation) = obj {
            let level = relation.tags
                                .get("admin_level")
                                .and_then(|s| s.parse().ok());
            let level = match level {
                None => {
                    info!("invalid admin_level for relation {}: admin_level {:?}",
                          relation.id,
                          relation.tags.get("admin_level"));
                    continue;
                }
                Some(l) => l,
            };
            // administrative region with name ?
            let name = match relation.tags.get("name") {
                Some(val) => val,
                None => {
                    warn!("adminstrative region without name for relation {}:  admin_level {} \
                           ignored.",
                          relation.id,
                          level);
                    continue;
                }
            };
            // admininstrative region without coordinates
            let coord_centre = relation.refs
                                       .iter()
                                       .find(|rf| rf.role == "admin_centre")
                                       .and_then(|r| {
                                           objects.get(&r.member).and_then(|value| {
                                               match value {
                                                   &osmpbfreader::OsmObj::Node(ref node) => {
                                                       Some(mimir::Coord {
                                                           lat: node.lat,
                                                           lon: node.lon,
                                                       })
                                                   }
                                                   _ => None,
                                               }
                                           })
                                       });

            let admin_id = match relation.tags.get("ref:INSEE") {
                Some(val) => format!("admin:fr:{}", val.trim_left_matches('0')),
                None => format!("admin:osm:{}", relation.id),
            };
            let zip_code = match relation.tags.get("addr:postcode") {
                Some(val) => &val[..],
                None => "",
            };
            let boundary = build_boundary(&relation, &objects);
            let admin = mimir::Admin {
                id: admin_id,
                level: level,
                name: name.to_string(),
                zip_code: zip_code.to_string(),
                // TODO weight value ?
                weight: 1,
                coord: coord_centre,
                boundary: boundary,
            };
            administrative_regions.push(admin);
        }
    }
    return administrative_regions;
}

fn streets(pbf: &mut ParsedPbf) -> StreetsVec {

    let is_valid_street = |obj: &osmpbfreader::OsmObj| -> bool {
        match *obj {
            osmpbfreader::OsmObj::Way(ref way) => {
                way.tags.get("highway").map_or(false, |x: &String| !x.is_empty()) &&
                way.tags.get("name").map_or(false, |x: &String| !x.is_empty())
            }
            _ => false,
        }
    };

    let is_associated_street_relation = |obj: &osmpbfreader::OsmObj| -> bool {
        match *obj {
            osmpbfreader::OsmObj::Relation(ref rel) => {
                rel.tags.get("type").map_or(false, |v| v == "associatedStreet")
            }
            _ => false,
        }
    };

    let mut street_map = osmpbfreader::get_objs_and_deps(pbf, is_valid_street).unwrap();
    let relation_map = osmpbfreader::get_objs_and_deps(pbf, is_associated_street_relation).unwrap();

    // Sometimes, streets can be devided into several "way"s that still have the same street name.
    // The reason why a street is devided may be that a part of the street become a bridge/tunne/etc.
    // In this case, a "relation" tagged with (type = associatedStreet) is used to group all these "way"s.
    // In order not to have doublons in autocompleion, we should keep only one "way" in the relation
    // and remove all the rest way whose street name is the same.
    for (_, rel_obj) in &relation_map {
        if let &osmpbfreader::OsmObj::Relation(ref rel) = rel_obj {
            let mut found = false;
            for ref_obj in &rel.refs {
                if let &osmpbfreader::OsmId::Way(_) = &ref_obj.member {
                    if !found {
                        found = true;
                        continue;
                    };
                    street_map.remove(&ref_obj.member);
                };
            }
        };
    }

    street_map.iter()
              .filter_map(|(osm_id, obj)| {
                  if let &osmpbfreader::OsmObj::Way(ref way) = obj {
                      if let &osmpbfreader::OsmId::Way(ref way_id) = osm_id {
                          way.tags.get("name").and_then(|way_name| {
                              Some(mimirsbrunn::Street {
                                  id: way_id.to_string(),
                                  street_name: way_name.to_string(),
                                  name: way_name.to_string(),
                                  weight: 1,
                                  administrative_region: None,
                              })
                          })
                      } else {
                          None
                      }
                  } else {
                      None
                  }
              })
              .collect()

}

fn index_osm(es_cnx_string: &str, admins: &AdminsVec, streets: &StreetsVec) {
    let mut rubber = Rubber::new(es_cnx_string);
    rubber.create_index();
    match rubber.clean_db_by_doc_type(&["admin", "street"]) {
        Err(e) => panic!("failed to clean data by document type: {}", e),
        Ok(nb) => info!("clean data by document type : {}", nb),
    }
    info!("Add data in elasticsearch db.");
    match rubber.bulk_index(admins.iter()) {
        Err(e) => panic!("failed to index admins of osm because: {}", e),
        Ok(nb) => info!("Nb of indexed adminstrative regions: {}", nb),
    }

    match rubber.bulk_index(streets.iter()) {
        Err(e) => panic!("failed to index streets of osm because: {}", e),
        Ok(nb) => info!("Nb of indexed streets: {}", nb),
    }
}

fn main() {
    mimir::logger_init().unwrap();
    debug!("importing adminstrative region into Mimir");
    let args: Args = docopt::Docopt::new(USAGE)
                         .and_then(|d| d.decode())
                         .unwrap_or_else(|e| e.exit());

    let levels = args.flag_level.iter().cloned().collect();
    let mut parsed_pbf = parse_osm_pbf(&args.flag_input);
    let admins = administrative_regions(&mut parsed_pbf, levels);
    let streets = streets(&mut parsed_pbf);
    index_osm(&args.flag_connection_string, &admins, &streets);
}
<|endoftext|>


<commit_before>#![allow(unused_imports)]
#![allow(non_snake_case)]

use std::cell::RefCell;
use std::cmp::{max, min, Ordering};
use std::collections::*;
use std::fmt::{Debug, Formatter, Write as FmtWrite};
use std::io::{stderr, stdin, BufRead, Write};
use std::mem::{replace, swap};
use std::ops::*;
use std::rc::Rc;

// -----------------------------------------------
// Framework
// -----------------------------------------------

#[allow(unused_macros)]
macro_rules! read {
    ([$t:ty] ; $n:expr) =>
        ((0..$n).map(|_| read!([$t])).collect::<Vec<_>>());
    ($($t:ty),+ ; $n:expr) =>
        ((0..$n).map(|_| read!($($t),+)).collect::<Vec<_>>());
    ([$t:ty]) =>
        (rl().split_whitespace().map(|w| w.parse().unwrap()).collect::<Vec<$t>>());
    ($t:ty) =>
        (rl().parse::<$t>().unwrap());
    ($($t:ty),*) => {{
        let buf = rl();
        let mut w = buf.split_whitespace();
        ($(w.next().unwrap().parse::<$t>().unwrap()),*)
    }};
}

#[allow(unused_macros)]
macro_rules! debug {
    ($($arg:expr),*) => {
        #[cfg(debug_assertions)]
        $(writeln!(stderr(), "{} = {:?}", stringify!($arg), $arg).unwrap());*
    };
}

#[allow(dead_code)]
fn rl() -> String {
    let mut buf = String::new();
    stdin().read_line(&mut buf).unwrap();
    buf.trim_right().to_owned()
}

trait IteratorExt: Iterator + Sized {
    fn vec(self) -> Vec<Self::Item> {
        self.collect()
    }
}

impl<T: Iterator> IteratorExt for T {}

// -----------------------------------------------
// Solution
// -----------------------------------------------

fn main() {
    println!("{}", 0)
}
<commit_msg>Add GitHub link to main<commit_after>#![allow(unused_imports)]
#![allow(non_snake_case)]

use std::cell::RefCell;
use std::cmp::{max, min, Ordering};
use std::collections::*;
use std::fmt::{Debug, Formatter, Write as FmtWrite};
use std::io::{stderr, stdin, BufRead, Write};
use std::mem::{replace, swap};
use std::ops::*;
use std::rc::Rc;

// -----------------------------------------------
// Framework <https://github.com/vain0x/procon>
// -----------------------------------------------

#[allow(unused_macros)]
macro_rules! read {
    ([$t:ty] ; $n:expr) =>
        ((0..$n).map(|_| read!([$t])).collect::<Vec<_>>());
    ($($t:ty),+ ; $n:expr) =>
        ((0..$n).map(|_| read!($($t),+)).collect::<Vec<_>>());
    ([$t:ty]) =>
        (rl().split_whitespace().map(|w| w.parse().unwrap()).collect::<Vec<$t>>());
    ($t:ty) =>
        (rl().parse::<$t>().unwrap());
    ($($t:ty),*) => {{
        let buf = rl();
        let mut w = buf.split_whitespace();
        ($(w.next().unwrap().parse::<$t>().unwrap()),*)
    }};
}

#[allow(unused_macros)]
macro_rules! debug {
    ($($arg:expr),*) => {
        #[cfg(debug_assertions)]
        $(writeln!(stderr(), "{} = {:?}", stringify!($arg), $arg).unwrap());*
    };
}

#[allow(dead_code)]
fn rl() -> String {
    let mut buf = String::new();
    stdin().read_line(&mut buf).unwrap();
    buf.trim_right().to_owned()
}

trait IteratorExt: Iterator + Sized {
    fn vec(self) -> Vec<Self::Item> {
        self.collect()
    }
}

impl<T: Iterator> IteratorExt for T {}

// -----------------------------------------------
// Solution
// -----------------------------------------------

fn main() {
    println!("{}", 0)
}
<|endoftext|>


<commit_before>use error::{CliError, Usage, UsageKind};
use super::Command;

extern crate chrono;

use self::chrono::Local;

use error::CannotHappenError;
use lib::fs::*;
use lib::io::*;
use lib::setting;
use std::fs;
use std::io::Error as IoError;
use std::path::PathBuf;

pub struct Sweep;

impl Command for Sweep {
    fn validation(&self) -> bool { true }

    fn usage(&self) -> Usage {
        return Usage::new(UsageKind::Sweep);
    }

    fn main(&self) -> Result<(), CliError> {
        let now  = Local::now();
        let date = now.format("%Y-%m-%d").to_string();

        let path_to_dust_box = path_buf![setting::storage_dir(), date, "dusts"];
        try!(setting::create_essential_dir_all(&path_to_dust_box));

        let ignore = try!(setting::read_ignore_file());

        let target_files = walk_dir(".")
            .difference(&ignore)
            .cloned()
            .collect::<Vec<String>>();
        for target in &target_files {
            try!(Self::move_file_to_dust_box(target, &path_to_dust_box));
        }

        let dirs = dirs_ordered_by_descending_depth(".");
        for target in &dirs {
            if is_empty_dir(target) {
                try!(Self::move_dir_to_dust_box(target, &path_to_dust_box));
            }
        }

        Ok(())
    }
}

impl Sweep {
    fn move_file_to_dust_box(target: &String, path_to_dust_box: &PathBuf) -> Result<(), CliError> {
        let target_path = path_buf![target];
        let target_file = try!(target_path.file_name().ok_or(CannotHappenError));
        let target_base = try!(target_path.parent().ok_or(CannotHappenError));
        let to = path_buf![&path_to_dust_box, target_base];

        let message = format!("Move \"{}\" file to \"{}\" directory", target_path.display(), path_to_dust_box.display());
        print_with_tag(Tag::Info, message);

        try!(setting::create_essential_dir_all(&to));

        // forcedly overwrite if the file exists with same name.
        try!(fs::rename(target, path_buf![to, target_file]));

        Ok(())
    }

    fn move_dir_to_dust_box(target: &String, path_to_dust_box: &PathBuf) -> Result<(), IoError> {
        try!(fs::remove_dir(target));
        try!(setting::create_essential_dir_all(&path_buf![path_to_dust_box, target]));

        Ok(())
    }
}
<commit_msg>[Remedy] Ignore permission error on "sweep"<commit_after>use error::{CliError, Usage, UsageKind};
use super::Command;

extern crate chrono;

use self::chrono::Local;

use error::CannotHappenError;
use lib::fs::*;
use lib::io::*;
use lib::setting;
use std::fs;
use std::io::{Error as IoError, ErrorKind as IoErrorKind};
use std::path::PathBuf;

pub struct Sweep;

impl Command for Sweep {
    fn validation(&self) -> bool { true }

    fn usage(&self) -> Usage {
        return Usage::new(UsageKind::Sweep);
    }

    fn main(&self) -> Result<(), CliError> {
        let now  = Local::now();
        let date = now.format("%Y-%m-%d").to_string();

        let path_to_dust_box = path_buf![setting::storage_dir(), date, "dusts"];
        try!(setting::create_essential_dir_all(&path_to_dust_box));

        let ignore = try!(setting::read_ignore_file());

        let target_files = walk_dir(".")
            .difference(&ignore)
            .cloned()
            .collect::<Vec<String>>();
        for target in &target_files {
            try!(Self::move_file_to_dust_box(target, &path_to_dust_box));
        }

        let all_dirs = dirs_ordered_by_descending_depth(".");
        for target in &all_dirs {
            try!(Self::move_empty_dir_to_dust_box(target, &path_to_dust_box));
        }

        Ok(())
    }
}

impl Sweep {
    fn move_file_to_dust_box(target: &String, path_to_dust_box: &PathBuf) -> Result<(), CliError> {
        let target_path = path_buf![target];
        let target_file = try!(target_path.file_name().ok_or(CannotHappenError));
        let target_base = try!(target_path.parent().ok_or(CannotHappenError));
        let to = path_buf![&path_to_dust_box, target_base];

        let message = format!("Move file \"{}\" to \"{}\" directory", target_path.display(), path_to_dust_box.display());
        print_with_tag(Tag::Info, message);

        try!(setting::create_essential_dir_all(&to));

        // forcedly overwrite if the file exists with same name.
        match fs::rename(target, path_buf![to, target_file]) {
            Ok(_)  => (),
            Err(e) => match e.kind() {
                IoErrorKind::PermissionDenied => print_with_tag(Tag::Info, "Interrupted this moving file for permission"),
                _                             => return Err(From::from(e)),
            },
        };

        Ok(())
    }

    fn move_empty_dir_to_dust_box(target: &String, path_to_dust_box: &PathBuf) -> Result<(), IoError> {
        if is_empty_dir(target) {
            let message = format!("Move empty directory \"{}\" to \"{}\" directory", target, path_to_dust_box.display());
            print_with_tag(Tag::Info, message);

            match fs::remove_dir(target) {
                Ok(_)  => try!(setting::create_essential_dir_all(&path_buf![path_to_dust_box, target])),
                Err(e) => match e.kind() {
                    IoErrorKind::PermissionDenied => print_with_tag(Tag::Info, "Interrupted this moving directory for permission"),
                    _                             => return Err(From::from(e)),
                },
            };
        }

        Ok(())
    }
}
<|endoftext|>


<commit_before>
use std;

/// A unique identifier.
#[derive(Copy,Clone,Debug)]
pub struct Id(usize);

/// A special value signifying an unspeficied Id.
const UNSPECIFIED_ID: usize = !0;

impl Id
{
    /// Gets an ID which represents
    pub fn unspecified() -> Self {
        Id(UNSPECIFIED_ID)
    }

    /// Checks if the ID is unspecified.
    pub fn is_specified(self) -> bool {
        let Id(val) = self;
        val != UNSPECIFIED_ID
    }

    /// Gets the underlying ID.
    pub fn underlying(self) -> usize {
        let Id(val) = self;
        val
    }
}

/// An object which has an identifier.
pub trait Identifiable
{
    /// Sets the internal ID of the object.
    /// This **should not** be called manually.
    fn set_id(&mut self, id: Id);
}

/// A unique identifier generator.
#[derive(Copy,Clone,Debug)]
pub struct Generator
{
    next: usize,
}

impl Generator
{
    /// Creates a new generator.
    pub fn new() -> Self {
        Generator {
            next: 0,
        }
    }

    /// Gets a new identifier.
    pub fn next(&mut self) -> Id {

        if self.next == UNSPECIFIED_ID {
            panic!("the ID space has been exhausted");
        }

        let id = self.next;
        self.next += 1;

        Id(id)
    }
}

impl std::cmp::PartialEq for Id
{
    fn eq(&self, other: &Id) -> bool {

        assert!(self.is_specified() && other.is_specified(),
                "id's have not been assigned");

        let &Id(i1) = self;
        let &Id(i2) = other;

        i1 == i2
    }
}

impl std::cmp::Eq for Id { }

impl std::fmt::Display for Id
{
    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
        let &Id(val) = self;
        val.fmt(fmt)
    }
}
<commit_msg>Added check to Id<commit_after>
use std;

/// A unique identifier.
#[derive(Copy,Clone,Debug)]
pub struct Id(usize);

/// A special value signifying an unspeficied Id.
const UNSPECIFIED_ID: usize = !0;

impl Id
{
    /// Gets an ID which represents
    pub fn unspecified() -> Self {
        Id(UNSPECIFIED_ID)
    }

    /// Checks if the ID is unspecified.
    pub fn is_specified(self) -> bool {
        let Id(val) = self;
        val != UNSPECIFIED_ID
    }

    /// Gets the underlying ID.
    pub fn underlying(self) -> usize {
        let Id(val) = self;
        val
    }
}

/// An object which has an identifier.
pub trait Identifiable
{
    /// Sets the internal ID of the object.
    /// This **should not** be called manually.
    fn set_id(&mut self, id: Id);
}

/// A unique identifier generator.
#[derive(Copy,Clone,Debug)]
pub struct Generator
{
    next: usize,
}

impl Generator
{
    /// Creates a new generator.
    pub fn new() -> Self {
        Generator {
            next: 0,
        }
    }

    /// Gets a new identifier.
    pub fn next(&mut self) -> Id {

        if self.next == UNSPECIFIED_ID {
            panic!("the ID space has been exhausted");
        }

        let id = self.next;
        self.next += 1;

        Id(id)
    }
}

impl std::cmp::PartialEq for Id
{
    fn eq(&self, other: &Id) -> bool {

        assert!(self.is_specified() && other.is_specified(),
                "id's have not been assigned");

        let &Id(i1) = self;
        let &Id(i2) = other;

        i1 == i2
    }
}

impl std::cmp::Eq for Id { }

impl std::fmt::Display for Id
{
    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
        assert!(self.is_specified(),
                "ID is not set");

        let &Id(val) = self;
        val.fmt(fmt)
    }
}
<|endoftext|>


<commit_before>use super::{LzfResult,LzfError};
use std::cmp;

const HLOG    : usize = 16;
const HSIZE   : u32 = 1 << HLOG;
const MAX_OFF : usize = 1 << 13;
const MAX_REF : usize = ((1 << 8) + (1 << 3));
const MAX_LIT : i32 = (1 << 5);

#[inline(always)]
fn get(d: &[u8], i: usize) -> u8 {
    unsafe { *d.get_unchecked(i) }
}

#[inline(always)]
fn first(p: &[u8], off: usize) -> u32 {
    ((get(p,off) as u32) << 8) | get(p,off+1) as u32
}

#[inline(always)]
fn next(v: u32, p: &[u8], off: usize) -> u32 {
    (v << 8) | get(p,off+2) as u32
}

#[inline(always)]
fn idx(h: u32) -> usize {
    let h = h as u64;
    (
        // 8 = 3*8-HLOG, but HLOG is constant at 16
        (h.wrapping_shr(8).wrapping_sub(h*5))
        & (HSIZE-1) as u64
    ) as usize
}

#[inline(always)]
fn not(i: i32) -> i32 {
    if i == 0 { 1 } else { 0 }
}

pub fn compress(data: &[u8]) -> LzfResult<Vec<u8>> {
    let in_len = data.len();
    let out_buf_len = in_len;
    let mut out = Vec::with_capacity(out_buf_len);
    unsafe { out.set_len(out_buf_len) };

    let mut out_len : i32 = 1; /* start run by default */

    let mut htab = [0; 1<<HLOG];

    let mut current_offset = 0;

    if in_len == 0 {
        return Err(LzfError::DataCorrupted);
    }

    let mut lit : i32 = 0;

    let mut hval : u32;
    let mut ref_offset;

    hval = first(data, current_offset);

    while current_offset < in_len-2 {
        hval = next(hval, data, current_offset);
        let hslot_idx = idx(hval);

        unsafe {
            ref_offset = *htab.get_unchecked(hslot_idx);
            *htab.get_unchecked_mut(hslot_idx) = current_offset;
        }

        let off = current_offset.wrapping_sub(ref_offset).wrapping_sub(1);
        if off < MAX_OFF
            && current_offset+4 < in_len
            && ref_offset > 0
            && get(data,ref_offset+0) == get(data,current_offset+0)
            && get(data,ref_offset+1) == get(data,current_offset+1)
            && get(data,ref_offset+2) == get(data,current_offset+2) {

            let mut len = 2;
            let maxlen = cmp::min(in_len - current_offset - len, MAX_REF);

            /* stop run */
            unsafe { *out.get_unchecked_mut((out_len - lit - 1) as usize) = (lit as u8).wrapping_sub(1); }
            out_len -= not(lit); /* undo run if length is zero */

            if out_len as i32 + 3 + 1 >= out_buf_len as i32 {
                return Err(LzfError::NoCompressionPossible);
            }

            loop {
                // Unrool loop.
                if maxlen > 16 {
                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }
                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }
                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }
                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }

                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }
                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }
                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }
                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }

                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }
                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }
                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }
                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }

                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }
                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }
                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }
                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }
                }

                len += 1;
                while len < maxlen && get(data,ref_offset+len) == get(data,current_offset+len) {
                    len += 1;
                }
                break;
            }

            len -= 2; /* len is now #octets - 1 */
            current_offset += 1;

            if len < 7 {
                unsafe{ *out.get_unchecked_mut(out_len as usize) = (off >> 8) as u8 + (len << 5) as u8; }
                out_len += 1;
            } else {
                unsafe {
                    *out.get_unchecked_mut(out_len as usize) = (off >> 8) as u8 + (7 << 5);
                    *out.get_unchecked_mut(out_len as usize + 1) = (len as u8).wrapping_sub(7);
                }
                out_len += 2;
            }

            unsafe { *out.get_unchecked_mut(out_len as usize) = off as u8; }
            out_len += 2; /* start run */
            lit = 0;

            /* we add here, because we later substract from the total length */
            current_offset += len-1;

            if current_offset >= in_len {
                break;
            }

            hval = first(data, current_offset);

            hval = next(hval, data, current_offset);
            unsafe { *htab.get_unchecked_mut(idx(hval)) = current_offset; }
            current_offset += 1;

            hval = next(hval, data, current_offset);
            unsafe { *htab.get_unchecked_mut(idx(hval)) = current_offset; }
            current_offset += 1;
        } else {
            /* one more literal byte we must copy */
            if current_offset >= out_buf_len {
                return Err(LzfError::NoCompressionPossible);
            }

            lit += 1;
            unsafe { *out.get_unchecked_mut(out_len as usize) = get(data,current_offset); }
            out_len += 1;
            current_offset += 1;

            if lit == MAX_LIT {
                /* stop run */
                unsafe { *out.get_unchecked_mut((out_len - lit - 1) as usize) = (lit as u8).wrapping_sub(1); }
                lit = 0;
                out_len += 1; /* start run */
            }
        }
    }

    /* at most 3 bytes can be missing here */
    if out_len + 3 > out_buf_len as i32 {
        return Err(LzfError::NoCompressionPossible);
    }

    while current_offset < in_len {
        lit += 1;
        unsafe { *out.get_unchecked_mut(out_len as usize) = get(data,current_offset); }
        out_len += 1;
        current_offset += 1;

        if lit == MAX_LIT {
            /* stop run */
            unsafe { *out.get_unchecked_mut((out_len - lit - 1) as usize) = (lit as u8).wrapping_sub(1); }
            lit = 0;
            out_len += 1; /* start run */
        }
    }

    /* end run */
    unsafe { *out.get_unchecked_mut((out_len - lit - 1) as usize) = (lit as u8).wrapping_sub(1); }
    out_len -= not(lit); /* undo run if length is zero */

    unsafe { out.set_len(out_len as usize) };

    Ok(out)
}

#[test]
fn test_compress_skips_short() {
    match compress("foo".as_bytes()) {
        Ok(_) => panic!("Compression did _something_, which is wrong for 'foo'"),
        Err(err) => assert_eq!(LzfError::NoCompressionPossible, err)
    }
}

#[test]
fn test_compress_lorem() {
    let lorem = "Lorem ipsum dolor sit amet, consetetur sadipscing elitr, sed diam nonumy eirmod tempor invidunt ut labore et dolore magna aliquyam erat, sed diam voluptua. At vero eos et accusam et justo duo dolores et ea rebum. Stet clita kasd gubergren, no sea takimata sanctus est Lorem ipsum dolor sit amet. Lorem ipsum dolor sit amet, consetetur sadipscing elitr, sed diam nonumy eirmod tempor invidunt ut labore et dolore magna aliquyam erat, sed diam voluptua.";

    match compress(lorem.as_bytes()) {
        Ok(compressed) => {
            assert_eq!(272, compressed.len())
        }
        Err(err) => panic!("Compression failed with error {:?}", err)
    }
}

#[test]
fn test_compress_decompress_lorem_round() {
    use super::native;

    let lorem = "Lorem ipsum dolor sit amet, consetetur sadipscing elitr, sed diam nonumy eirmod tempor invidunt ut labore et dolore magna aliquyam erat, sed diam voluptua. At vero eos et accusam et justo duo dolores et ea rebum. Stet clita kasd gubergren, no sea takimata sanctus est Lorem ipsum dolor sit amet. Lorem ipsum dolor sit amet, consetetur sadipscing elitr, sed diam nonumy eirmod tempor invidunt ut labore et dolore magna aliquyam erat, sed diam voluptua.";

    let compressed = match compress(lorem.as_bytes()) {
        Ok(c) => c,
        Err(err) => panic!("Compression failed with error {:?}", err)
    };

    match native::decompress(&compressed, lorem.len()) {
        Ok(decompressed) => {
            assert_eq!(lorem.len(), decompressed.len());
            assert_eq!(lorem.as_bytes(), &decompressed[..]);
        },
        Err(err) => panic!("Decompression failed with error {:?}", err)
    };
}

#[test]
fn test_alice_wonderland_both() {
    let alice = "\r\n\r\n\r\n\r\n                ALICE'S ADVENTURES IN WONDERLAND\r\n";

    let compressed = match compress(alice.as_bytes()) {
        Ok(c) => c,
        Err(err) => panic!("Compression failed with error {:?}", err)
    };

    let c_compressed = match super::compress(alice.as_bytes()) {
        Ok(c) => c,
        Err(err) => panic!("Compression failed with error {:?}", err)
    };

    assert_eq!(&compressed[..], &c_compressed[..]);
}
<commit_msg>Don't force inlining, the compiler probably knows better<commit_after>use super::{LzfResult,LzfError};
use std::cmp;

const HLOG    : usize = 16;
const HSIZE   : u32 = 1 << HLOG;
const MAX_OFF : usize = 1 << 13;
const MAX_REF : usize = ((1 << 8) + (1 << 3));
const MAX_LIT : i32 = (1 << 5);

fn get(d: &[u8], i: usize) -> u8 {
    unsafe { *d.get_unchecked(i) }
}

fn first(p: &[u8], off: usize) -> u32 {
    ((get(p,off) as u32) << 8) | get(p,off+1) as u32
}

fn next(v: u32, p: &[u8], off: usize) -> u32 {
    (v << 8) | get(p,off+2) as u32
}

fn idx(h: u32) -> usize {
    let h = h as u64;
    (
        // 8 = 3*8-HLOG, but HLOG is constant at 16
        (h.wrapping_shr(8).wrapping_sub(h*5))
        & (HSIZE-1) as u64
    ) as usize
}

fn not(i: i32) -> i32 {
    if i == 0 { 1 } else { 0 }
}

pub fn compress(data: &[u8]) -> LzfResult<Vec<u8>> {
    let in_len = data.len();
    let out_buf_len = in_len;
    let mut out = Vec::with_capacity(out_buf_len);
    unsafe { out.set_len(out_buf_len) };

    let mut out_len : i32 = 1; /* start run by default */

    let mut htab = [0; 1<<HLOG];

    let mut current_offset = 0;

    if in_len == 0 {
        return Err(LzfError::DataCorrupted);
    }

    let mut lit : i32 = 0;

    let mut hval : u32;
    let mut ref_offset;

    hval = first(data, current_offset);

    while current_offset < in_len-2 {
        hval = next(hval, data, current_offset);
        let hslot_idx = idx(hval);

        unsafe {
            ref_offset = *htab.get_unchecked(hslot_idx);
            *htab.get_unchecked_mut(hslot_idx) = current_offset;
        }

        let off = current_offset.wrapping_sub(ref_offset).wrapping_sub(1);
        if off < MAX_OFF
            && current_offset+4 < in_len
            && ref_offset > 0
            && get(data,ref_offset+0) == get(data,current_offset+0)
            && get(data,ref_offset+1) == get(data,current_offset+1)
            && get(data,ref_offset+2) == get(data,current_offset+2) {

            let mut len = 2;
            let maxlen = cmp::min(in_len - current_offset - len, MAX_REF);

            /* stop run */
            unsafe { *out.get_unchecked_mut((out_len - lit - 1) as usize) = (lit as u8).wrapping_sub(1); }
            out_len -= not(lit); /* undo run if length is zero */

            if out_len as i32 + 3 + 1 >= out_buf_len as i32 {
                return Err(LzfError::NoCompressionPossible);
            }

            loop {
                // Unrool loop.
                if maxlen > 16 {
                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }
                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }
                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }
                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }

                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }
                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }
                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }
                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }

                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }
                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }
                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }
                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }

                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }
                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }
                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }
                    len += 1; if get(data,ref_offset+len) != get(data,current_offset+len) { break; }
                }

                len += 1;
                while len < maxlen && get(data,ref_offset+len) == get(data,current_offset+len) {
                    len += 1;
                }
                break;
            }

            len -= 2; /* len is now #octets - 1 */
            current_offset += 1;

            if len < 7 {
                unsafe{ *out.get_unchecked_mut(out_len as usize) = (off >> 8) as u8 + (len << 5) as u8; }
                out_len += 1;
            } else {
                unsafe {
                    *out.get_unchecked_mut(out_len as usize) = (off >> 8) as u8 + (7 << 5);
                    *out.get_unchecked_mut(out_len as usize + 1) = (len as u8).wrapping_sub(7);
                }
                out_len += 2;
            }

            unsafe { *out.get_unchecked_mut(out_len as usize) = off as u8; }
            out_len += 2; /* start run */
            lit = 0;

            /* we add here, because we later substract from the total length */
            current_offset += len-1;

            if current_offset >= in_len {
                break;
            }

            hval = first(data, current_offset);

            hval = next(hval, data, current_offset);
            unsafe { *htab.get_unchecked_mut(idx(hval)) = current_offset; }
            current_offset += 1;

            hval = next(hval, data, current_offset);
            unsafe { *htab.get_unchecked_mut(idx(hval)) = current_offset; }
            current_offset += 1;
        } else {
            /* one more literal byte we must copy */
            if current_offset >= out_buf_len {
                return Err(LzfError::NoCompressionPossible);
            }

            lit += 1;
            unsafe { *out.get_unchecked_mut(out_len as usize) = get(data,current_offset); }
            out_len += 1;
            current_offset += 1;

            if lit == MAX_LIT {
                /* stop run */
                unsafe { *out.get_unchecked_mut((out_len - lit - 1) as usize) = (lit as u8).wrapping_sub(1); }
                lit = 0;
                out_len += 1; /* start run */
            }
        }
    }

    /* at most 3 bytes can be missing here */
    if out_len + 3 > out_buf_len as i32 {
        return Err(LzfError::NoCompressionPossible);
    }

    while current_offset < in_len {
        lit += 1;
        unsafe { *out.get_unchecked_mut(out_len as usize) = get(data,current_offset); }
        out_len += 1;
        current_offset += 1;

        if lit == MAX_LIT {
            /* stop run */
            unsafe { *out.get_unchecked_mut((out_len - lit - 1) as usize) = (lit as u8).wrapping_sub(1); }
            lit = 0;
            out_len += 1; /* start run */
        }
    }

    /* end run */
    unsafe { *out.get_unchecked_mut((out_len - lit - 1) as usize) = (lit as u8).wrapping_sub(1); }
    out_len -= not(lit); /* undo run if length is zero */

    unsafe { out.set_len(out_len as usize) };

    Ok(out)
}

#[test]
fn test_compress_skips_short() {
    match compress("foo".as_bytes()) {
        Ok(_) => panic!("Compression did _something_, which is wrong for 'foo'"),
        Err(err) => assert_eq!(LzfError::NoCompressionPossible, err)
    }
}

#[test]
fn test_compress_lorem() {
    let lorem = "Lorem ipsum dolor sit amet, consetetur sadipscing elitr, sed diam nonumy eirmod tempor invidunt ut labore et dolore magna aliquyam erat, sed diam voluptua. At vero eos et accusam et justo duo dolores et ea rebum. Stet clita kasd gubergren, no sea takimata sanctus est Lorem ipsum dolor sit amet. Lorem ipsum dolor sit amet, consetetur sadipscing elitr, sed diam nonumy eirmod tempor invidunt ut labore et dolore magna aliquyam erat, sed diam voluptua.";

    match compress(lorem.as_bytes()) {
        Ok(compressed) => {
            assert_eq!(272, compressed.len())
        }
        Err(err) => panic!("Compression failed with error {:?}", err)
    }
}

#[test]
fn test_compress_decompress_lorem_round() {
    use super::native;

    let lorem = "Lorem ipsum dolor sit amet, consetetur sadipscing elitr, sed diam nonumy eirmod tempor invidunt ut labore et dolore magna aliquyam erat, sed diam voluptua. At vero eos et accusam et justo duo dolores et ea rebum. Stet clita kasd gubergren, no sea takimata sanctus est Lorem ipsum dolor sit amet. Lorem ipsum dolor sit amet, consetetur sadipscing elitr, sed diam nonumy eirmod tempor invidunt ut labore et dolore magna aliquyam erat, sed diam voluptua.";

    let compressed = match compress(lorem.as_bytes()) {
        Ok(c) => c,
        Err(err) => panic!("Compression failed with error {:?}", err)
    };

    match native::decompress(&compressed, lorem.len()) {
        Ok(decompressed) => {
            assert_eq!(lorem.len(), decompressed.len());
            assert_eq!(lorem.as_bytes(), &decompressed[..]);
        },
        Err(err) => panic!("Decompression failed with error {:?}", err)
    };
}

#[test]
fn test_alice_wonderland_both() {
    let alice = "\r\n\r\n\r\n\r\n                ALICE'S ADVENTURES IN WONDERLAND\r\n";

    let compressed = match compress(alice.as_bytes()) {
        Ok(c) => c,
        Err(err) => panic!("Compression failed with error {:?}", err)
    };

    let c_compressed = match super::compress(alice.as_bytes()) {
        Ok(c) => c,
        Err(err) => panic!("Compression failed with error {:?}", err)
    };

    assert_eq!(&compressed[..], &c_compressed[..]);
}
<|endoftext|>


<commit_before>use std::{cmp, ptr};
use std::io::{self,Write,Read};
use std::iter::repeat;
use pool::Reset;

pub struct Buffer {
  memory:   Vec<u8>,
  capacity: usize,
  position: usize,
  end:      usize
}

impl Buffer {
  pub fn with_capacity(capacity: usize) -> Buffer {
    let mut v = Vec::with_capacity(capacity);
    v.extend(repeat(0).take(capacity));
    Buffer {
      memory:   v,
      capacity: capacity,
      position: 0,
      end:      0
    }
  }

  pub fn available_data(&self) -> usize {
    self.end - self.position
  }

  pub fn available_space(&self) -> usize {
    self.capacity - self.end
  }

  pub fn consume(&mut self, count: usize) -> usize {
    let cnt        = cmp::min(count, self.available_data());
    self.position += cnt;
    if self.position > self.capacity / 2 {
      //trace!("consume shift: pos {}, end {}", self.position, self.end);
      self.shift();
    }
    cnt
  }

  pub fn fill(&mut self, count: usize) -> usize {
    let cnt   = cmp::min(count, self.available_space());
    self.end += cnt;
    if self.available_space() < self.available_data() + cnt {
      //trace!("fill shift: pos {}, end {}", self.position, self.end);
      self.shift();
    }

    cnt
  }

  pub fn reset(&mut self) {
    self.position = 0;
    self.end      = 0;
  }

  pub fn data(&self) -> &[u8] {
    &self.memory[self.position..self.end]
  }

  pub fn space(&mut self) -> &mut[u8] {
    &mut self.memory[self.end..self.capacity]
  }

  pub fn shift(&mut self) {
    if self.position > 0 {
      unsafe {
        let length = self.end - self.position;
        ptr::copy( (&self.memory[self.position..self.end]).as_ptr(), (&mut self.memory[..length]).as_mut_ptr(), length);
        self.position = 0;
        self.end      = length;
      }
    }
  }

  pub fn delete_slice(&mut self, start: usize, length: usize) -> Option<usize> {
    if start + length >= self.available_data() {
      return None
    }

    unsafe {
      let begin    = self.position + start;
      let next_end = self.end - length;
      ptr::copy(
        (&self.memory[begin+length..self.end]).as_ptr(),
        (&mut self.memory[begin..next_end]).as_mut_ptr(),
        self.end - (begin+length)
      );
      self.end = next_end;
    }
    Some(self.available_data())
  }

  pub fn replace_slice(&mut self, data: &[u8], start: usize, length: usize) -> Option<usize> {
    let data_len = data.len();
    if start + length > self.available_data() ||
      self.position + start + data_len > self.capacity {
      return None
    }

    unsafe {
      let begin     = self.position + start;
      let slice_end = begin + data_len;
      // we reduced the data size
      if data_len < length {
        ptr::copy(data.as_ptr(), (&mut self.memory[begin..slice_end]).as_mut_ptr(), data_len);

        ptr::copy((&self.memory[start+length..self.end]).as_ptr(), (&mut self.memory[slice_end..]).as_mut_ptr(), self.end - (start + length));
        self.end = self.end - (length - data_len);

      // we put more data in the buffer
      } else {
        ptr::copy((&self.memory[start+length..self.end]).as_ptr(), (&mut self.memory[start+data_len..]).as_mut_ptr(), self.end - (start + length));
        ptr::copy(data.as_ptr(), (&mut self.memory[begin..slice_end]).as_mut_ptr(), data_len);
        self.end = self.end + data_len - length;
      }
    }
    Some(self.available_data())
  }
}

impl Write for Buffer {
  fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
    match self.space().write(buf) {
      Ok(size) => { self.fill(size); Ok(size) },
      err      => err
    }
  }

  fn flush(&mut self) -> io::Result<()> {
    Ok(())
  }
}

impl Read for Buffer {
  fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
    let len = cmp::min(self.available_data(), buf.len());
    unsafe {
      ptr::copy((&self.memory[self.position..len]).as_ptr(), buf.as_mut_ptr(), len);
      self.position += len;
    }
    Ok(len)
  }
}

impl Reset for Buffer {
  fn reset(&mut self) {
    self.reset();
  }
}

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

  #[test]
  fn fill_and_consume() {
    let mut b = Buffer::with_capacity(10);
    assert_eq!(b.available_data(), 0);
    assert_eq!(b.available_space(), 10);
    let res = b.write(&b"abcd"[..]);
    assert_eq!(res.ok(), Some(4));
    assert_eq!(b.available_data(), 4);
    assert_eq!(b.available_space(), 6);

    assert_eq!(b.data(), &b"abcd"[..]);

    b.consume(2);
    assert_eq!(b.available_data(), 2);
    assert_eq!(b.available_space(), 6);
    assert_eq!(b.data(), &b"cd"[..]);

    b.shift();
    assert_eq!(b.available_data(), 2);
    assert_eq!(b.available_space(), 8);
    assert_eq!(b.data(), &b"cd"[..]);

    assert_eq!(b.write(&b"efghijklmnop"[..]).ok(), Some(8));
    assert_eq!(b.available_data(), 10);
    assert_eq!(b.available_space(), 0);
    assert_eq!(b.data(), &b"cdefghijkl"[..]);
    b.shift();
    assert_eq!(b.available_data(), 10);
    assert_eq!(b.available_space(), 0);
    assert_eq!(b.data(), &b"cdefghijkl"[..]);
  }

  #[test]
  fn delete() {
    let mut b = Buffer::with_capacity(10);
    let res = b.write(&b"abcdefgh"[..]);
    assert_eq!(b.available_data(), 8);
    assert_eq!(b.available_space(), 2);

    assert_eq!(b.delete_slice(2, 3), Some(5));
    assert_eq!(b.available_data(), 5);
    assert_eq!(b.available_space(), 5);
    assert_eq!(b.data(), &b"abfgh"[..]);

    assert_eq!(b.delete_slice(5, 2), None);
    assert_eq!(b.delete_slice(4, 2), None);
  }

  #[test]
  fn replace() {
    let mut b = Buffer::with_capacity(10);
    let res = b.write(&b"abcdefgh"[..]);
    assert_eq!(b.available_data(), 8);
    assert_eq!(b.available_space(), 2);

    assert_eq!(b.replace_slice(&b"ABC"[..], 2, 3), Some(8));
    assert_eq!(b.available_data(), 8);
    assert_eq!(b.available_space(), 2);
    assert_eq!(b.data(), &b"abABCfgh"[..]);

    assert_eq!(b.replace_slice(&b"XYZ"[..], 8, 3), None);
    assert_eq!(b.replace_slice(&b"XYZ"[..], 6, 3), None);

    assert_eq!(b.replace_slice(&b"XYZ"[..], 2, 4), Some(7));
    assert_eq!(b.available_data(), 7);
    assert_eq!(b.available_space(), 3);
    assert_eq!(b.data(), &b"abXYZgh"[..]);

    assert_eq!(b.replace_slice(&b"123"[..], 2, 2), Some(8));
    assert_eq!(b.available_data(), 8);
    assert_eq!(b.available_space(), 2);
    assert_eq!(b.data(), &b"ab123Zgh"[..]);
  }
}
<commit_msg>Implement insert_slice for buffers<commit_after>use std::{cmp, ptr};
use std::io::{self,Write,Read};
use std::iter::repeat;
use pool::Reset;

pub struct Buffer {
  memory:   Vec<u8>,
  capacity: usize,
  position: usize,
  end:      usize
}

impl Buffer {
  pub fn with_capacity(capacity: usize) -> Buffer {
    let mut v = Vec::with_capacity(capacity);
    v.extend(repeat(0).take(capacity));
    Buffer {
      memory:   v,
      capacity: capacity,
      position: 0,
      end:      0
    }
  }

  pub fn available_data(&self) -> usize {
    self.end - self.position
  }

  pub fn available_space(&self) -> usize {
    self.capacity - self.end
  }

  pub fn consume(&mut self, count: usize) -> usize {
    let cnt        = cmp::min(count, self.available_data());
    self.position += cnt;
    if self.position > self.capacity / 2 {
      //trace!("consume shift: pos {}, end {}", self.position, self.end);
      self.shift();
    }
    cnt
  }

  pub fn fill(&mut self, count: usize) -> usize {
    let cnt   = cmp::min(count, self.available_space());
    self.end += cnt;
    if self.available_space() < self.available_data() + cnt {
      //trace!("fill shift: pos {}, end {}", self.position, self.end);
      self.shift();
    }

    cnt
  }

  pub fn reset(&mut self) {
    self.position = 0;
    self.end      = 0;
  }

  pub fn data(&self) -> &[u8] {
    &self.memory[self.position..self.end]
  }

  pub fn space(&mut self) -> &mut[u8] {
    &mut self.memory[self.end..self.capacity]
  }

  pub fn shift(&mut self) {
    if self.position > 0 {
      unsafe {
        let length = self.end - self.position;
        ptr::copy( (&self.memory[self.position..self.end]).as_ptr(), (&mut self.memory[..length]).as_mut_ptr(), length);
        self.position = 0;
        self.end      = length;
      }
    }
  }

  pub fn delete_slice(&mut self, start: usize, length: usize) -> Option<usize> {
    if start + length >= self.available_data() {
      return None
    }

    unsafe {
      let begin    = self.position + start;
      let next_end = self.end - length;
      ptr::copy(
        (&self.memory[begin+length..self.end]).as_ptr(),
        (&mut self.memory[begin..next_end]).as_mut_ptr(),
        self.end - (begin+length)
      );
      self.end = next_end;
    }
    Some(self.available_data())
  }

  pub fn replace_slice(&mut self, data: &[u8], start: usize, length: usize) -> Option<usize> {
    let data_len = data.len();
    if start + length > self.available_data() ||
      self.position + start + data_len > self.capacity {
      return None
    }

    unsafe {
      let begin     = self.position + start;
      let slice_end = begin + data_len;
      // we reduced the data size
      if data_len < length {
        ptr::copy(data.as_ptr(), (&mut self.memory[begin..slice_end]).as_mut_ptr(), data_len);

        ptr::copy((&self.memory[start+length..self.end]).as_ptr(), (&mut self.memory[slice_end..]).as_mut_ptr(), self.end - (start + length));
        self.end = self.end - (length - data_len);

      // we put more data in the buffer
      } else {
        ptr::copy((&self.memory[start+length..self.end]).as_ptr(), (&mut self.memory[start+data_len..]).as_mut_ptr(), self.end - (start + length));
        ptr::copy(data.as_ptr(), (&mut self.memory[begin..slice_end]).as_mut_ptr(), data_len);
        self.end = self.end + data_len - length;
      }
    }
    Some(self.available_data())
  }

  pub fn insert_slice(&mut self, data: &[u8], start: usize) -> Option<usize> {
    let data_len = data.len();
    if start > self.available_data() ||
      self.position + start + data_len > self.capacity {
      return None
    }

    unsafe {
      let begin     = self.position + start;
      let slice_end = begin + data_len;
      ptr::copy((&self.memory[start..self.end]).as_ptr(), (&mut self.memory[start+data_len..]).as_mut_ptr(), self.end - start);
      ptr::copy(data.as_ptr(), (&mut self.memory[begin..slice_end]).as_mut_ptr(), data_len);
      self.end = self.end + data_len;
    }
    Some(self.available_data())
  }
}

impl Write for Buffer {
  fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
    match self.space().write(buf) {
      Ok(size) => { self.fill(size); Ok(size) },
      err      => err
    }
  }

  fn flush(&mut self) -> io::Result<()> {
    Ok(())
  }
}

impl Read for Buffer {
  fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
    let len = cmp::min(self.available_data(), buf.len());
    unsafe {
      ptr::copy((&self.memory[self.position..len]).as_ptr(), buf.as_mut_ptr(), len);
      self.position += len;
    }
    Ok(len)
  }
}

impl Reset for Buffer {
  fn reset(&mut self) {
    self.reset();
  }
}

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

  #[test]
  fn fill_and_consume() {
    let mut b = Buffer::with_capacity(10);
    assert_eq!(b.available_data(), 0);
    assert_eq!(b.available_space(), 10);
    let res = b.write(&b"abcd"[..]);
    assert_eq!(res.ok(), Some(4));
    assert_eq!(b.available_data(), 4);
    assert_eq!(b.available_space(), 6);

    assert_eq!(b.data(), &b"abcd"[..]);

    b.consume(2);
    assert_eq!(b.available_data(), 2);
    assert_eq!(b.available_space(), 6);
    assert_eq!(b.data(), &b"cd"[..]);

    b.shift();
    assert_eq!(b.available_data(), 2);
    assert_eq!(b.available_space(), 8);
    assert_eq!(b.data(), &b"cd"[..]);

    assert_eq!(b.write(&b"efghijklmnop"[..]).ok(), Some(8));
    assert_eq!(b.available_data(), 10);
    assert_eq!(b.available_space(), 0);
    assert_eq!(b.data(), &b"cdefghijkl"[..]);
    b.shift();
    assert_eq!(b.available_data(), 10);
    assert_eq!(b.available_space(), 0);
    assert_eq!(b.data(), &b"cdefghijkl"[..]);
  }

  #[test]
  fn delete() {
    let mut b = Buffer::with_capacity(10);
    let res = b.write(&b"abcdefgh"[..]);
    assert_eq!(b.available_data(), 8);
    assert_eq!(b.available_space(), 2);

    assert_eq!(b.delete_slice(2, 3), Some(5));
    assert_eq!(b.available_data(), 5);
    assert_eq!(b.available_space(), 5);
    assert_eq!(b.data(), &b"abfgh"[..]);

    assert_eq!(b.delete_slice(5, 2), None);
    assert_eq!(b.delete_slice(4, 2), None);
  }

  #[test]
  fn replace() {
    let mut b = Buffer::with_capacity(10);
    let res = b.write(&b"abcdefgh"[..]);
    assert_eq!(b.available_data(), 8);
    assert_eq!(b.available_space(), 2);

    assert_eq!(b.replace_slice(&b"ABC"[..], 2, 3), Some(8));
    assert_eq!(b.available_data(), 8);
    assert_eq!(b.available_space(), 2);
    assert_eq!(b.data(), &b"abABCfgh"[..]);

    assert_eq!(b.replace_slice(&b"XYZ"[..], 8, 3), None);
    assert_eq!(b.replace_slice(&b"XYZ"[..], 6, 3), None);

    assert_eq!(b.replace_slice(&b"XYZ"[..], 2, 4), Some(7));
    assert_eq!(b.available_data(), 7);
    assert_eq!(b.available_space(), 3);
    assert_eq!(b.data(), &b"abXYZgh"[..]);

    assert_eq!(b.replace_slice(&b"123"[..], 2, 2), Some(8));
    assert_eq!(b.available_data(), 8);
    assert_eq!(b.available_space(), 2);
    assert_eq!(b.data(), &b"ab123Zgh"[..]);
  }
}
<|endoftext|>