<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,
    flag_import_way: bool
}

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

Options:
    -h, --help            Show this message.
    -i, --input=<file>    OSM PBF file.
    -l, --level=<level>   Admin levels to keep.
    -w, --import-way      By default, streets are not imported
    -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: &str) -> 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| !x.is_empty()) &&
                way.tags.get("name").map_or(false, |x| !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()

}
// TODO: template these two functions
fn index_admin(es_cnx_string: &str, admins: AdminsVec) {
    let mut rubber = Rubber::new(es_cnx_string);
    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),
    }
}

fn index_ways(es_cnx_string: &str, streets: StreetsVec) {
    let mut rubber = Rubber::new(es_cnx_string);
    match rubber.bulk_index(streets.iter()) {
        Err(e) => panic!("failed to index streets of osm because: {}", e),
        Ok(nb) => info!("Nb of indexed street: {}", nb),
    }
}

fn index_settings(es_cnx_string: &str) {
    let mut rubber = Rubber::new(es_cnx_string);
    rubber.create_index();
}

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);

    index_settings(&args.flag_connection_string);
    index_admin(&args.flag_connection_string, admins);

    let mut optinoal_streets: Option<StreetsVec> = None;
    if args.flag_import_way  {
        optinoal_streets = Some(streets(&mut parsed_pbf));
    }

    if let Some(streets) = optinoal_streets {
	index_ways(&args.flag_connection_string, streets);
    }
}
<commit_msg>read pdf only once in street<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,
    flag_import_way: bool,
}

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

Options:
    -h, --help            Show this message.
    -i, --input=<file>    OSM PBF file.
    -l, --level=<level>   Admin levels to keep.
    -w, --import-way      By default, streets are not imported
    -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: &str) -> 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| !x.is_empty()) &&
                way.tags.get("name").map_or(false, |x| !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 objs_map = osmpbfreader::get_objs_and_deps(pbf, is_valid_street).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.
    let mut objs_to_remove = Vec::<osmpbfreader::OsmId>::new();

    for (_, rel_obj) in &objs_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;
                    };
                    objs_to_remove.push(ref_obj.member.clone());
                };
            }
        };
    }

    for osm_id in objs_to_remove {
        objs_map.remove(&osm_id);
    }

    objs_map.iter()
            .filter_map(|(_, obj)| {
                if let &osmpbfreader::OsmObj::Way(ref way) = obj {
                    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
                }
            })
            .collect()

}
// TODO: template these two functions
fn index_admin(es_cnx_string: &str, admins: AdminsVec) {
    let mut rubber = Rubber::new(es_cnx_string);

    match rubber.clean_db_by_doc_type(&["admin"]) {
        Err(e) => panic!("failed to clean data by document type: {}", e),
        Ok(nb) => info!("Nb of cleaned indexed obj: {}", nb),
    }

    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),
    }
}

fn index_ways(es_cnx_string: &str, streets: StreetsVec) {
    let mut rubber = Rubber::new(es_cnx_string);

    match rubber.clean_db_by_doc_type(&["street"]) {
        Err(e) => panic!("failed to clean data by document type: {}", e),
        Ok(nb) => info!("Nb of cleaned indexed obj: {}", nb),
    }

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

fn index_settings(es_cnx_string: &str) {
    let mut rubber = Rubber::new(es_cnx_string);
    rubber.create_index();
}

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);

    index_settings(&args.flag_connection_string);
    index_admin(&args.flag_connection_string, admins);

    let mut optinoal_streets: Option<StreetsVec> = None;
    if args.flag_import_way {
        optinoal_streets = Some(streets(&mut parsed_pbf));
    }

    if let Some(streets) = optinoal_streets {
        index_ways(&args.flag_connection_string, streets);
    }
}
<|endoftext|>


<commit_before>//! Functions to choose good boundaries for block splitting. Deflate allows encoding
//! the data in multiple blocks, with a separate Huffman tree for each block. The
//! Huffman tree itself requires some bytes to encode, so by choosing certain
//! blocks, you can either hurt, or enhance compression. These functions choose good
//! ones that enhance it.

use std::io::Write;
use std::mem::{transmute, uninitialized};
use std::ptr::null_mut;

use libc::{c_void, size_t};
use libc::funcs::c95::stdlib::{free, malloc};

use super::Options;

use deflate::calculate_block_size;
use lz77::{BlockState, LZ77Store, lz77_greedy, clean_lz77_store};
use util;

/// The "f" for the FindMinimum function below.
/// i: the current parameter of f(i)
/// context: for your implementation
type FindMinimumFun = unsafe fn(i: usize, context: *const c_void) -> f64;

/// Finds minimum of function f(i) where is is of type size_t, f(i) is of type
/// double, i is in range start-end (excluding end).
unsafe fn find_minimum(f: FindMinimumFun, context: *const c_void, start: usize, end: usize) -> usize {
    if end - start < 1024 {
        let mut best: f64 = util::LARGE_FLOAT;
        let mut result = start;
        for i in start..end {
            let v = f(i, context);
            if v < best {
                best = v;
                result = i;
            }
        }
        result
    } else {
        // Try to find minimum faster by recursively checking multiple points.

        // Good value: 9.
        const NUM: usize = 9;

        let mut p: [usize; NUM] = uninitialized();
        let mut vp: [f64; NUM] = uninitialized();
        let mut lastbest: f64 = util::LARGE_FLOAT;
        let mut pos = start;

        let mut start = start;
        let mut end = end;
        loop {
            if end - start <= NUM {
                break;
            }

            for i in 0..NUM {
                p[i] = start + (i + 1) * ((end - start) / (NUM + 1));
                vp[i] = f(p[i], context);
            }
            let mut besti: usize = 0;
            let mut best: f64 = vp[0];
            for i in 1..NUM {
                if vp[i] < best {
                    best = vp[i];
                    besti = i;
                }
            }
            if best > lastbest {
                break;
            }

            start = if besti == 0 { start } else { p[besti - 1] };
            end = if besti == NUM - 1 { end } else { p[besti + 1] };

            pos = p[besti];
            lastbest = best;
        }
        pos
    }
}

/**
 * Returns estimated cost of a block in bits.  It includes the size to encode the
 * tree and the size to encode all literal, length and distance symbols and their
 * extra bits.
 *
 * litlens: lz77 lit/lengths
 * dists: ll77 distances
 * lstart: start of block
 * lend: end of block (not inclusive)
 */
unsafe fn estimate_cost(litlens: *const u16, dists: *const u16, lstart: usize, lend: usize) -> f64 {
    calculate_block_size(litlens, dists, lstart, lend, 2)
}

struct SplitCostContext {
    litlens: *const u16,
    dists: *const u16,
    _llsize: usize,
    start: usize,
    end: usize,
}

/// Gets the cost which is the sum of the cost of the left and the right section
/// of the data.
/// type: FindMinimumFun
unsafe fn split_cost(i: usize, context: *const c_void) -> f64 {
    let c: *const SplitCostContext = transmute(context);
    estimate_cost((*c).litlens, (*c).dists, (*c).start, i) + estimate_cost((*c).litlens, (*c).dists, i, (*c).end)
}

unsafe fn add_sorted(value: usize, out: *mut *mut usize, outsize: *mut usize) {
    append_data!(value, *out, *outsize);
    let mut i: usize = 0;
    while i + 1 < *outsize {
        if *(*out).offset(i as isize) > value {
            let mut j: usize = *outsize - 1;
            while j > i {
                *(*out).offset(j as isize) = *(*out).offset(j as isize - 1);
                j -= 1;
            }
            *(*out).offset(i as isize) = value;
            break;
        }
        i += 1;
    }
}

/// Prints the block split points as decimal and hex values in the terminal.
unsafe fn print_block_split_points(litlens: *const u16, dists: *const u16, llsize: usize, lz77splitpoints: *const usize, nlz77points: usize) {
    let mut splitpoints: *mut usize = null_mut();
    let mut npoints: usize = 0;

    // The input is given as lz77 indices, but we want to see the uncompressed
    // index values.
    let mut pos: usize = 0;
    if nlz77points > 0 {
        for i in 0..llsize {
            let length: usize = if *dists.offset(i as isize) == 0 { 1 } else { *litlens.offset(i as isize) as usize };
            if *lz77splitpoints.offset(npoints as isize) == i {
                append_data!(pos, splitpoints, npoints);
                if npoints == nlz77points {
                    break;
                }
            }
            pos += length;
        }
    }
    assert_eq!(npoints, nlz77points);

    print_err!("block split points: ");
    for i in 0..npoints {
        print_err!("{} ", *splitpoints.offset(i as isize) as i32);
    }
    print_err!("(hex:");
    for i in 0..npoints {
        print_err!(" {:x}", *splitpoints.offset(i as isize) as i32);
    }
    println_err!("");

    free(splitpoints as *mut c_void);
}

/// Finds next block to try to split, the largest of the available ones.
/// The largest is chosen to make sure that if only a limited amount of blocks is
/// requested, their sizes are spread evenly.
/// llsize: the size of the LL77 data, which is the size of the done array here.
/// done: array indicating which blocks starting at that position are no longer
///     splittable (splitting them increases rather than decreases cost).
/// splitpoints: the splitpoints found so far.
/// npoints: the amount of splitpoints found so far.
/// lstart: output variable, giving start of block.
/// lend: output variable, giving end of block.
/// returns 1 if a block was found, 0 if no block found (all are done).
unsafe fn find_largest_splittable_block(llsize: usize, done: *const u8, splitpoints: *const usize, npoints: usize, lstart: *mut usize, lend: *mut usize) -> bool {
    let mut longest: usize = 0;
    let mut found = false;
    for i in 0..npoints+1 {
        let start: usize = if i == 0 { 0 } else { *splitpoints.offset(i as isize - 1) };
        let end: usize = if i == npoints { llsize - 1 } else { *splitpoints.offset(i as isize) };
        if *done.offset(start as isize) == 0 && end - start > longest {
            *lstart = start;
            *lend = end;
            found = true;
            longest = end - start;
        }
    }
    found
}

/// Does blocksplitting on LZ77 data.
/// The output splitpoints are indices in the LZ77 data.
/// litlens: lz77 lit/lengths
/// dists: lz77 distances
/// llsize: size of litlens and dists
/// maxblocks: set a limit to the amount of blocks. Set to 0 to mean no limit.
pub unsafe fn block_split_lz77(options: *const Options, litlens: *const u16, dists: *const u16, llsize: usize, maxblocks: usize, splitpoints: *mut *mut usize, npoints: *mut usize) {
    if llsize < 10 {
        // This code fails on tiny files.
        return;
    }

    let done: *mut u8 = transmute(malloc(llsize as size_t));
    if done.is_null() {
        panic!("Allocation failed!");
    }
    for i in 0..llsize {
        *done.offset(i as isize) = 0;
    }

    let mut lstart: usize = 0;
    let mut lend: usize = llsize;
    let mut numblocks: usize = 1;
    loop {
        if maxblocks > 0 && numblocks >= maxblocks {
            break;
        }

        let mut c = SplitCostContext {
            litlens: litlens,
            dists: dists,
            _llsize: llsize,
            start: lstart,
            end: lend,
        };
        assert!(lstart < lend);
        let llpos: usize = find_minimum(split_cost, &mut c as *mut _ as *mut c_void, lstart + 1, lend);

        assert!(llpos > lstart);
        assert!(llpos < lend);

        let splitcost: f64 = estimate_cost(litlens, dists, lstart, llpos) + estimate_cost(litlens, dists, llpos, lend);
        let origcost: f64 = estimate_cost(litlens, dists, lstart, lend);

        if splitcost > origcost || llpos == lstart + 1 || llpos == lend {
            *done.offset(lstart as isize) = 1;
        } else {
            add_sorted(llpos, splitpoints, npoints);
            numblocks += 1;
        }

        if !find_largest_splittable_block(llsize, done, *splitpoints, *npoints, &mut lstart, &mut lend) {
            // No further split will probably reduce compression.
            break;
        }

        if lend - lstart < 10 {
            break;
        }
    }

    if (*options).verbose {
        print_block_split_points(litlens, dists, llsize, *splitpoints, *npoints);
    }

    free(transmute(done));
}

/**
 * Does blocksplitting on uncompressed data.
 * The output splitpoints are indices in the uncompressed bytes.
 *
 * options: general program options.
 * in: uncompressed input data
 * instart: where to start splitting
 * inend: where to end splitting (not inclusive)
 * maxblocks: maximum amount of blocks to split into, or 0 for no limit
 * splitpoints: dynamic array to put the resulting split point coordinates into.
 *   The coordinates are indices in the input array.
 * npoints: pointer to amount of splitpoints, for the dynamic array. The amount of
 *   blocks is the amount of splitpoitns + 1.
 */
pub unsafe fn block_split(options: *const Options, in_: *const u8, instart: usize, inend: usize, maxblocks: usize, splitpoints: *mut *mut usize, npoints: *mut usize) {
    let s = BlockState::new(options, instart, inend, null_mut());
    let mut lz77splitpoints: *mut usize = null_mut();
    let mut nlz77points: usize = 0;
    let mut store = LZ77Store::new();

    *npoints = 0;
    *splitpoints = null_mut();

    // Unintuitively, Using a simple LZ77 method here instead of ZopfliLZ77Optimal
    // results in better blocks.
    lz77_greedy(&s, in_, instart, inend, &mut store);

    block_split_lz77(options, store.litlens, store.dists, store.size, maxblocks, &mut lz77splitpoints, &mut nlz77points);

    // Convert LZ77 positions to positions in the uncompressed input.
    let mut pos: usize = instart;
    if nlz77points > 0 {
        for i in 0..store.size {
            let length: usize = if *store.dists.offset(i as isize) == 0 { 1 } else { *store.litlens.offset(i as isize) as usize };
            if *lz77splitpoints.offset(*npoints as isize) == i {
                append_data!(pos, *splitpoints, *npoints);
                if *npoints == nlz77points {
                    break;
                }
            }
            pos += length;
        }
    }
    assert_eq!(*npoints, nlz77points);

    free(lz77splitpoints as *mut c_void);
    clean_lz77_store(&mut store);
}

/// Divides the input into equal blocks, does not even take LZ77 lengths into
/// account.
pub unsafe fn block_split_simple(_in: *const u8, instart: usize, inend: usize, blocksize: usize, splitpoints: *mut *mut usize, npoints: *mut usize) {
    let mut i: usize = instart;
    while i < inend {
        append_data!(i, *splitpoints, *npoints);
        i += blocksize;
    }
}
<commit_msg>Fix verbose output format<commit_after>//! Functions to choose good boundaries for block splitting. Deflate allows encoding
//! the data in multiple blocks, with a separate Huffman tree for each block. The
//! Huffman tree itself requires some bytes to encode, so by choosing certain
//! blocks, you can either hurt, or enhance compression. These functions choose good
//! ones that enhance it.

use std::io::Write;
use std::mem::{transmute, uninitialized};
use std::ptr::null_mut;

use libc::{c_void, size_t};
use libc::funcs::c95::stdlib::{free, malloc};

use super::Options;

use deflate::calculate_block_size;
use lz77::{BlockState, LZ77Store, lz77_greedy, clean_lz77_store};
use util;

/// The "f" for the FindMinimum function below.
/// i: the current parameter of f(i)
/// context: for your implementation
type FindMinimumFun = unsafe fn(i: usize, context: *const c_void) -> f64;

/// Finds minimum of function f(i) where is is of type size_t, f(i) is of type
/// double, i is in range start-end (excluding end).
unsafe fn find_minimum(f: FindMinimumFun, context: *const c_void, start: usize, end: usize) -> usize {
    if end - start < 1024 {
        let mut best: f64 = util::LARGE_FLOAT;
        let mut result = start;
        for i in start..end {
            let v = f(i, context);
            if v < best {
                best = v;
                result = i;
            }
        }
        result
    } else {
        // Try to find minimum faster by recursively checking multiple points.

        // Good value: 9.
        const NUM: usize = 9;

        let mut p: [usize; NUM] = uninitialized();
        let mut vp: [f64; NUM] = uninitialized();
        let mut lastbest: f64 = util::LARGE_FLOAT;
        let mut pos = start;

        let mut start = start;
        let mut end = end;
        loop {
            if end - start <= NUM {
                break;
            }

            for i in 0..NUM {
                p[i] = start + (i + 1) * ((end - start) / (NUM + 1));
                vp[i] = f(p[i], context);
            }
            let mut besti: usize = 0;
            let mut best: f64 = vp[0];
            for i in 1..NUM {
                if vp[i] < best {
                    best = vp[i];
                    besti = i;
                }
            }
            if best > lastbest {
                break;
            }

            start = if besti == 0 { start } else { p[besti - 1] };
            end = if besti == NUM - 1 { end } else { p[besti + 1] };

            pos = p[besti];
            lastbest = best;
        }
        pos
    }
}

/**
 * Returns estimated cost of a block in bits.  It includes the size to encode the
 * tree and the size to encode all literal, length and distance symbols and their
 * extra bits.
 *
 * litlens: lz77 lit/lengths
 * dists: ll77 distances
 * lstart: start of block
 * lend: end of block (not inclusive)
 */
unsafe fn estimate_cost(litlens: *const u16, dists: *const u16, lstart: usize, lend: usize) -> f64 {
    calculate_block_size(litlens, dists, lstart, lend, 2)
}

struct SplitCostContext {
    litlens: *const u16,
    dists: *const u16,
    _llsize: usize,
    start: usize,
    end: usize,
}

/// Gets the cost which is the sum of the cost of the left and the right section
/// of the data.
/// type: FindMinimumFun
unsafe fn split_cost(i: usize, context: *const c_void) -> f64 {
    let c: *const SplitCostContext = transmute(context);
    estimate_cost((*c).litlens, (*c).dists, (*c).start, i) + estimate_cost((*c).litlens, (*c).dists, i, (*c).end)
}

unsafe fn add_sorted(value: usize, out: *mut *mut usize, outsize: *mut usize) {
    append_data!(value, *out, *outsize);
    let mut i: usize = 0;
    while i + 1 < *outsize {
        if *(*out).offset(i as isize) > value {
            let mut j: usize = *outsize - 1;
            while j > i {
                *(*out).offset(j as isize) = *(*out).offset(j as isize - 1);
                j -= 1;
            }
            *(*out).offset(i as isize) = value;
            break;
        }
        i += 1;
    }
}

/// Prints the block split points as decimal and hex values in the terminal.
unsafe fn print_block_split_points(litlens: *const u16, dists: *const u16, llsize: usize, lz77splitpoints: *const usize, nlz77points: usize) {
    let mut splitpoints: *mut usize = null_mut();
    let mut npoints: usize = 0;

    // The input is given as lz77 indices, but we want to see the uncompressed
    // index values.
    let mut pos: usize = 0;
    if nlz77points > 0 {
        for i in 0..llsize {
            let length: usize = if *dists.offset(i as isize) == 0 { 1 } else { *litlens.offset(i as isize) as usize };
            if *lz77splitpoints.offset(npoints as isize) == i {
                append_data!(pos, splitpoints, npoints);
                if npoints == nlz77points {
                    break;
                }
            }
            pos += length;
        }
    }
    assert_eq!(npoints, nlz77points);

    print_err!("block split points: ");
    for i in 0..npoints {
        print_err!("{} ", *splitpoints.offset(i as isize) as i32);
    }
    print_err!("(hex:");
    for i in 0..npoints {
        print_err!(" {:x}", *splitpoints.offset(i as isize) as i32);
    }
    println_err!(")");

    free(splitpoints as *mut c_void);
}

/// Finds next block to try to split, the largest of the available ones.
/// The largest is chosen to make sure that if only a limited amount of blocks is
/// requested, their sizes are spread evenly.
/// llsize: the size of the LL77 data, which is the size of the done array here.
/// done: array indicating which blocks starting at that position are no longer
///     splittable (splitting them increases rather than decreases cost).
/// splitpoints: the splitpoints found so far.
/// npoints: the amount of splitpoints found so far.
/// lstart: output variable, giving start of block.
/// lend: output variable, giving end of block.
/// returns 1 if a block was found, 0 if no block found (all are done).
unsafe fn find_largest_splittable_block(llsize: usize, done: *const u8, splitpoints: *const usize, npoints: usize, lstart: *mut usize, lend: *mut usize) -> bool {
    let mut longest: usize = 0;
    let mut found = false;
    for i in 0..npoints+1 {
        let start: usize = if i == 0 { 0 } else { *splitpoints.offset(i as isize - 1) };
        let end: usize = if i == npoints { llsize - 1 } else { *splitpoints.offset(i as isize) };
        if *done.offset(start as isize) == 0 && end - start > longest {
            *lstart = start;
            *lend = end;
            found = true;
            longest = end - start;
        }
    }
    found
}

/// Does blocksplitting on LZ77 data.
/// The output splitpoints are indices in the LZ77 data.
/// litlens: lz77 lit/lengths
/// dists: lz77 distances
/// llsize: size of litlens and dists
/// maxblocks: set a limit to the amount of blocks. Set to 0 to mean no limit.
pub unsafe fn block_split_lz77(options: *const Options, litlens: *const u16, dists: *const u16, llsize: usize, maxblocks: usize, splitpoints: *mut *mut usize, npoints: *mut usize) {
    if llsize < 10 {
        // This code fails on tiny files.
        return;
    }

    let done: *mut u8 = transmute(malloc(llsize as size_t));
    if done.is_null() {
        panic!("Allocation failed!");
    }
    for i in 0..llsize {
        *done.offset(i as isize) = 0;
    }

    let mut lstart: usize = 0;
    let mut lend: usize = llsize;
    let mut numblocks: usize = 1;
    loop {
        if maxblocks > 0 && numblocks >= maxblocks {
            break;
        }

        let mut c = SplitCostContext {
            litlens: litlens,
            dists: dists,
            _llsize: llsize,
            start: lstart,
            end: lend,
        };
        assert!(lstart < lend);
        let llpos: usize = find_minimum(split_cost, &mut c as *mut _ as *mut c_void, lstart + 1, lend);

        assert!(llpos > lstart);
        assert!(llpos < lend);

        let splitcost: f64 = estimate_cost(litlens, dists, lstart, llpos) + estimate_cost(litlens, dists, llpos, lend);
        let origcost: f64 = estimate_cost(litlens, dists, lstart, lend);

        if splitcost > origcost || llpos == lstart + 1 || llpos == lend {
            *done.offset(lstart as isize) = 1;
        } else {
            add_sorted(llpos, splitpoints, npoints);
            numblocks += 1;
        }

        if !find_largest_splittable_block(llsize, done, *splitpoints, *npoints, &mut lstart, &mut lend) {
            // No further split will probably reduce compression.
            break;
        }

        if lend - lstart < 10 {
            break;
        }
    }

    if (*options).verbose {
        print_block_split_points(litlens, dists, llsize, *splitpoints, *npoints);
    }

    free(transmute(done));
}

/**
 * Does blocksplitting on uncompressed data.
 * The output splitpoints are indices in the uncompressed bytes.
 *
 * options: general program options.
 * in: uncompressed input data
 * instart: where to start splitting
 * inend: where to end splitting (not inclusive)
 * maxblocks: maximum amount of blocks to split into, or 0 for no limit
 * splitpoints: dynamic array to put the resulting split point coordinates into.
 *   The coordinates are indices in the input array.
 * npoints: pointer to amount of splitpoints, for the dynamic array. The amount of
 *   blocks is the amount of splitpoitns + 1.
 */
pub unsafe fn block_split(options: *const Options, in_: *const u8, instart: usize, inend: usize, maxblocks: usize, splitpoints: *mut *mut usize, npoints: *mut usize) {
    let s = BlockState::new(options, instart, inend, null_mut());
    let mut lz77splitpoints: *mut usize = null_mut();
    let mut nlz77points: usize = 0;
    let mut store = LZ77Store::new();

    *npoints = 0;
    *splitpoints = null_mut();

    // Unintuitively, Using a simple LZ77 method here instead of ZopfliLZ77Optimal
    // results in better blocks.
    lz77_greedy(&s, in_, instart, inend, &mut store);

    block_split_lz77(options, store.litlens, store.dists, store.size, maxblocks, &mut lz77splitpoints, &mut nlz77points);

    // Convert LZ77 positions to positions in the uncompressed input.
    let mut pos: usize = instart;
    if nlz77points > 0 {
        for i in 0..store.size {
            let length: usize = if *store.dists.offset(i as isize) == 0 { 1 } else { *store.litlens.offset(i as isize) as usize };
            if *lz77splitpoints.offset(*npoints as isize) == i {
                append_data!(pos, *splitpoints, *npoints);
                if *npoints == nlz77points {
                    break;
                }
            }
            pos += length;
        }
    }
    assert_eq!(*npoints, nlz77points);

    free(lz77splitpoints as *mut c_void);
    clean_lz77_store(&mut store);
}

/// Divides the input into equal blocks, does not even take LZ77 lengths into
/// account.
pub unsafe fn block_split_simple(_in: *const u8, instart: usize, inend: usize, blocksize: usize, splitpoints: *mut *mut usize, npoints: *mut usize) {
    let mut i: usize = instart;
    while i < inend {
        append_data!(i, *splitpoints, *npoints);
        i += blocksize;
    }
}
<|endoftext|>


<commit_before>use ansi_term::{ANSIString, Colour};
use config::Config;

/// Merges in the default values for the program
pub fn merge_defaults(c: &mut Config) {
    c.set_default("global.modules",
                     vec!["exit_code", "directory", "git", "prompt"])
        .unwrap();
    c.set_default("global.foreground", "bright_white").unwrap();
    c.set_default("global.background", "blue").unwrap();
    c.set_default("global.separator", "").unwrap();
    c.set_default("global.padding_left", " ").unwrap();
    c.set_default("global.padding_right", " ").unwrap();

    c.set_default("modules.directory.max_depth", 6).unwrap();

    c.set_default("modules.exit_code.bg_success", "green").unwrap();
    c.set_default("modules.exit_code.bg_error", "red").unwrap();

    // c.set_default("modules.git.symbol_clean", "").unwrap();
    c.set_default("modules.git.symbol_insertion", "+").unwrap();
    c.set_default("modules.git.symbol_deletion", "-").unwrap();
    // c.set_default("modules.git.symbol_push", "⇡").unwrap();
    // c.set_default("modules.git.symbol_pull", "⇣").unwrap();
    c.set_default("modules.git.show_diff_stats", true).unwrap();

    c.set_default("modules.prompt.output", "$").unwrap();
    c.set_default("modules.prompt.bg_success", "green").unwrap();
    c.set_default("modules.prompt.bg_error", "red").unwrap();
}

/// Generic formatting method
pub fn format_module<'a>(c: &mut Config,
                         name: &'a str,
                         output: Option<String>,
                         last_successful: Option<&str>)
                         -> (Option<&'a str>, Option<ANSIString<'static>>) {
    // Formatting was not successful if there was nothing to format
    if c.get_str(&format!("modules.{}.output", name)).is_none() && output.is_none() {
        return (None, None);
    }

    // Get config options
    let fg = c.get_str(&format!("modules.{}.foreground", name))
        .unwrap_or_else(|| c.get_str("global.foreground").unwrap_or_default());
    let fg = string_to_colour(fg);

    let bg = c.get_str(&format!("modules.{}.background", name))
        .unwrap_or_else(|| c.get_str("global.background").unwrap_or_default());
    let bg = string_to_colour(bg);

    let padding_left = c.get_str(&format!("modules.{}.padding_left", name))
        .unwrap_or_else(|| c.get_str("global.padding_left").unwrap_or_default());
    let padding_right = c.get_str(&format!("modules.{}.padding_right", name))
        .unwrap_or_else(|| c.get_str("global.padding_right").unwrap_or_default());
    let separator = c.get_str(&format!("modules.{}.separator", name))
        .unwrap_or_else(|| c.get_str("global.separator").unwrap_or_default());

    // Override the output if there is "output" in this module's config
    let output = if let Some(s) = c.get_str(&format!("modules.{}.output", name)) {
        s
    } else {
        output.unwrap()
    };

    // Format the main content
    let mut content = format!("\\[{}\\]{}{}{}\\[{}\\]",
                              fg.on(bg).prefix(),
                              padding_left,
                              output,
                              padding_right,
                              fg.on(bg).suffix());

    // Format the separator
    let fg = bg;
    if let Some(name) = last_successful {
        // There is a visible module that comes after this one
        let next_bg = c.get_str(&format!("modules.{}.background", name))
            .unwrap_or_else(|| c.get_str("global.background").unwrap_or_default());
        let next_bg = string_to_colour(next_bg);

        content = format!("{}\\[{}\\]{}\\[{}\\]",
                          content,
                          fg.on(next_bg).prefix(),
                          separator,
                          fg.on(next_bg).suffix());
    } else {
        // This is the final module
        content = format!("{}\\[{}\\]{}\\[{}\\]",
                          content,
                          fg.prefix(),
                          separator,
                          fg.suffix());
    }

    (Some(name), Some(ANSIString::from(content)))
}

// Converts a string (from the config file) to a Colour
// See: https://upload.wikimedia.org/wikipedia/commons/1/15/Xterm_256color_chart.svg
fn string_to_colour(s: String) -> Colour {
    if let Ok(i) = s.parse::<u8>() {
        return Colour::Fixed(i);
    }

    let s = s.to_lowercase();
    match s.as_ref() {
        "black" => Colour::Fixed(0),
        "bright_black" => Colour::Fixed(8),
        "red" => Colour::Fixed(1),
        "bright_red" => Colour::Fixed(9),
        "green" => Colour::Fixed(2),
        "bright_green" => Colour::Fixed(10),
        "yellow" => Colour::Fixed(3),
        "bright_yellow" => Colour::Fixed(11),
        "blue" => Colour::Fixed(4),
        "bright_blue" => Colour::Fixed(12),
        "purple" => Colour::Fixed(5),
        "bright_purple" => Colour::Fixed(13),
        "cyan" => Colour::Fixed(6),
        "bright_cyan" => Colour::Fixed(14),
        "white" => Colour::Fixed(7),
        "bright_white" => Colour::Fixed(15),
        _ => panic!("Invalid color option: {} in config file!", s),
    }
}

pub fn format_module_prompt<'a>(c: &mut Config,
                                last_successful: Option<&'a str>,
                                exit_code: &str)
                                -> (Option<&'a str>, Option<ANSIString<'static>>) {
    let bg = match exit_code.as_ref() {
        "0" => c.get_str("modules.prompt.bg_success").unwrap_or_default(),
        _ => c.get_str("modules.prompt.bg_error").unwrap_or_default(),
    };
    c.set("modules.prompt.background", bg).unwrap();

    let output = c.get_str("modules.prompt.output").unwrap_or_default();

    format_module(c, "prompt", Some(output), last_successful)
}

pub fn format_module_exit_code<'a>(c: &mut Config,
                                   last_successful: Option<&'a str>,
                                   exit_code: &str)
                                   -> (Option<&'a str>, Option<ANSIString<'static>>) {
    let bg = match exit_code.as_ref() {
        "0" => c.get_str("modules.exit_code.bg_success").unwrap_or_default(),
        _ => c.get_str("modules.exit_code.bg_error").unwrap_or_default(),
    };
    c.set("modules.exit_code.background", bg).unwrap();

    format_module(c, "exit_code", Some(exit_code.to_string()), last_successful)
}

pub fn format_module_directory<'a>(c: &mut Config,
                                   last_successful: Option<&'a str>)
                                   -> (Option<&'a str>, Option<ANSIString<'static>>) {
    use std::env;
    use std::path::PathBuf;

    let home = env::var("HOME").unwrap();
    let cwd = env::current_dir().unwrap();

    let mut directory = env::current_dir().unwrap();

    // Convert "/home/user/directory" to "~/directory"
    if let Ok(stripped_dir) = cwd.strip_prefix(&home) {
        directory = PathBuf::from("~").join(stripped_dir);
    }

    // Current directory depth
    let depth = directory.components().count();

    // Max number of directories we want to see
    let max_depth = c.get_int("modules.directory.max_depth").unwrap_or_default() as usize;

    // if depth > max_depth {
    //     let short_path = PathBuf::new();
    //     for i in 0..(max_depth / 2) {

    //     }
    // }

    format_module(c,
                  "directory",
                  Some(format!("{}", directory.display())),
                  last_successful)
}

pub fn format_module_git<'a>(c: &mut Config,
                             last_successful: Option<&'a str>)
                             -> (Option<&'a str>, Option<ANSIString<'static>>) {
    use git2::Repository;
    use std::env;

    let mut output = String::new();

    let repo = Repository::discover(env::current_dir().unwrap());
    if let Ok(repo) = repo {
        if let Ok(reference) = repo.head() {
            // Current branch name
            output.push_str(reference.shorthand().unwrap());

            let show_diffs = c.get_bool("modules.git.show_diff_stats").unwrap_or_default();
            // Insertions and deletions
            if show_diffs {
                if let Ok(diff) = repo.diff_index_to_workdir(None, None) {
                    if let Ok(stats) = diff.stats() {
                        if stats.files_changed() > 0 {
                            let symbol_deletion = c.get_str("modules.git.symbol_deletion")
                                .unwrap_or_default();
                            let symbol_insertion = c.get_str("modules.git.symbol_insertion")
                                .unwrap_or_default();

                            output.push_str(&format!(" ({}{}, {}{})",
                                                     symbol_deletion,
                                                     stats.deletions(),
                                                     symbol_insertion,
                                                     stats.insertions()));
                        }
                    }
                }
            }
        }
    }

    if output.is_empty() {
        return (None, None);
    }

    // let symbol_clean = c.get_str("modules.git.symbol_clean").unwrap_or_default();
    // let symbol_push = c.get_str("modules.git.symbol_push").unwrap_or_default();
    // let symbol_pull = c.get_str("modules.git.symbol_pull").unwrap_or_default();

    format_module(c, "git", Some(output), last_successful)
}

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

    #[test]
    fn test_string_to_colour() {
        const CONTENT: &'static str = "hi";

        struct Test<'a> {
            input: &'static str,
            expected: &'a str,
        }

        let tests = [Test {
                         input: "green",
                         expected: &format!("\x1B[38;5;2m{}\x1B[0m", CONTENT),
                     },
                     Test {
                         input: "bright_green",
                         expected: &format!("\x1B[38;5;10m{}\x1B[0m", CONTENT),
                     }];

        for test in &tests {
            let result = string_to_colour(test.input.to_string());
            let result = format!("{}", result.paint(CONTENT));

            assert_eq!(test.expected, result);
        }
    }

    #[test]
    #[should_panic]
    fn test_string_to_colour_invalid_input() {
        struct Test {
            input: &'static str,
        }

        let tests = [Test { input: "green" }, Test { input: "invalid" }];

        for test in &tests {
            string_to_colour(test.input.to_string());
        }
    }
}
<commit_msg>feat: contract long working directory paths<commit_after>use ansi_term::{ANSIString, Colour};
use config::Config;

/// Merges in the default values for the program
pub fn merge_defaults(c: &mut Config) {
    c.set_default("global.modules",
                     vec!["exit_code", "directory", "git", "prompt"])
        .unwrap();
    c.set_default("global.foreground", "bright_white").unwrap();
    c.set_default("global.background", "blue").unwrap();
    c.set_default("global.separator", "").unwrap();
    c.set_default("global.padding_left", " ").unwrap();
    c.set_default("global.padding_right", " ").unwrap();

    c.set_default("modules.directory.max_depth", 4).unwrap();

    c.set_default("modules.exit_code.bg_success", "green").unwrap();
    c.set_default("modules.exit_code.bg_error", "red").unwrap();

    // c.set_default("modules.git.symbol_clean", "").unwrap();
    c.set_default("modules.git.symbol_insertion", "+").unwrap();
    c.set_default("modules.git.symbol_deletion", "-").unwrap();
    // c.set_default("modules.git.symbol_push", "⇡").unwrap();
    // c.set_default("modules.git.symbol_pull", "⇣").unwrap();
    c.set_default("modules.git.show_diff_stats", true).unwrap();

    c.set_default("modules.prompt.output", "$").unwrap();
    c.set_default("modules.prompt.bg_success", "green").unwrap();
    c.set_default("modules.prompt.bg_error", "red").unwrap();
}

/// Generic formatting method
pub fn format_module<'a>(c: &mut Config,
                         name: &'a str,
                         output: Option<String>,
                         last_successful: Option<&str>)
                         -> (Option<&'a str>, Option<ANSIString<'static>>) {
    // Formatting was not successful if there was nothing to format
    if c.get_str(&format!("modules.{}.output", name)).is_none() && output.is_none() {
        return (None, None);
    }

    // Get config options
    let fg = c.get_str(&format!("modules.{}.foreground", name))
        .unwrap_or_else(|| c.get_str("global.foreground").unwrap_or_default());
    let fg = string_to_colour(fg);

    let bg = c.get_str(&format!("modules.{}.background", name))
        .unwrap_or_else(|| c.get_str("global.background").unwrap_or_default());
    let bg = string_to_colour(bg);

    let padding_left = c.get_str(&format!("modules.{}.padding_left", name))
        .unwrap_or_else(|| c.get_str("global.padding_left").unwrap_or_default());
    let padding_right = c.get_str(&format!("modules.{}.padding_right", name))
        .unwrap_or_else(|| c.get_str("global.padding_right").unwrap_or_default());
    let separator = c.get_str(&format!("modules.{}.separator", name))
        .unwrap_or_else(|| c.get_str("global.separator").unwrap_or_default());

    // Override the output if there is "output" in this module's config
    let output = if let Some(s) = c.get_str(&format!("modules.{}.output", name)) {
        s
    } else {
        output.unwrap()
    };

    // Format the main content
    let mut content = format!("\\[{}\\]{}{}{}\\[{}\\]",
                              fg.on(bg).prefix(),
                              padding_left,
                              output,
                              padding_right,
                              fg.on(bg).suffix());

    // Format the separator
    let fg = bg;
    if let Some(name) = last_successful {
        // There is a visible module that comes after this one
        let next_bg = c.get_str(&format!("modules.{}.background", name))
            .unwrap_or_else(|| c.get_str("global.background").unwrap_or_default());
        let next_bg = string_to_colour(next_bg);

        content = format!("{}\\[{}\\]{}\\[{}\\]",
                          content,
                          fg.on(next_bg).prefix(),
                          separator,
                          fg.on(next_bg).suffix());
    } else {
        // This is the final module
        content = format!("{}\\[{}\\]{}\\[{}\\]",
                          content,
                          fg.prefix(),
                          separator,
                          fg.suffix());
    }

    (Some(name), Some(ANSIString::from(content)))
}

// Converts a string (from the config file) to a Colour
// See: https://upload.wikimedia.org/wikipedia/commons/1/15/Xterm_256color_chart.svg
fn string_to_colour(s: String) -> Colour {
    if let Ok(i) = s.parse::<u8>() {
        return Colour::Fixed(i);
    }

    let s = s.to_lowercase();
    match s.as_ref() {
        "black" => Colour::Fixed(0),
        "bright_black" => Colour::Fixed(8),
        "red" => Colour::Fixed(1),
        "bright_red" => Colour::Fixed(9),
        "green" => Colour::Fixed(2),
        "bright_green" => Colour::Fixed(10),
        "yellow" => Colour::Fixed(3),
        "bright_yellow" => Colour::Fixed(11),
        "blue" => Colour::Fixed(4),
        "bright_blue" => Colour::Fixed(12),
        "purple" => Colour::Fixed(5),
        "bright_purple" => Colour::Fixed(13),
        "cyan" => Colour::Fixed(6),
        "bright_cyan" => Colour::Fixed(14),
        "white" => Colour::Fixed(7),
        "bright_white" => Colour::Fixed(15),
        _ => panic!("Invalid color option: {} in config file!", s),
    }
}

pub fn format_module_prompt<'a>(c: &mut Config,
                                last_successful: Option<&'a str>,
                                exit_code: &str)
                                -> (Option<&'a str>, Option<ANSIString<'static>>) {
    let bg = match exit_code.as_ref() {
        "0" => c.get_str("modules.prompt.bg_success").unwrap_or_default(),
        _ => c.get_str("modules.prompt.bg_error").unwrap_or_default(),
    };
    c.set("modules.prompt.background", bg).unwrap();

    let output = c.get_str("modules.prompt.output").unwrap_or_default();

    format_module(c, "prompt", Some(output), last_successful)
}

pub fn format_module_exit_code<'a>(c: &mut Config,
                                   last_successful: Option<&'a str>,
                                   exit_code: &str)
                                   -> (Option<&'a str>, Option<ANSIString<'static>>) {
    let bg = match exit_code.as_ref() {
        "0" => c.get_str("modules.exit_code.bg_success").unwrap_or_default(),
        _ => c.get_str("modules.exit_code.bg_error").unwrap_or_default(),
    };
    c.set("modules.exit_code.background", bg).unwrap();

    format_module(c, "exit_code", Some(exit_code.to_string()), last_successful)
}

pub fn format_module_directory<'a>(c: &mut Config,
                                   last_successful: Option<&'a str>)
                                   -> (Option<&'a str>, Option<ANSIString<'static>>) {
    use std::env;
    use std::path::PathBuf;

    let home = env::var("HOME").unwrap();
    let cwd = env::current_dir().unwrap();

    // Convert "/home/user/directory" to "~/directory"
    let mut shortened_cwd: PathBuf;
    if let Ok(stripped_cwd) = cwd.strip_prefix(&home) {
        shortened_cwd = PathBuf::from("~").join(stripped_cwd);
    } else {
        shortened_cwd = env::current_dir().unwrap();
    }

    let depth = shortened_cwd.components().count();

    // Max number of directories we want to see
    let max_depth = c.get_int("modules.directory.max_depth").unwrap_or_default() as usize;

    if depth > max_depth {
        let comp_iter = shortened_cwd.clone();
        let comp_iter = comp_iter.components();

        shortened_cwd = PathBuf::new();
        for (i, component) in comp_iter.enumerate() {
            if i < (max_depth / 2) || i >= (depth - (max_depth / 2)) {
                shortened_cwd.push(component.as_os_str());
            } else if i == (max_depth / 2) {
                shortened_cwd.push("...");
            }
        }
    }

    format_module(c,
                  "directory",
                  Some(format!("{}", shortened_cwd.display())),
                  last_successful)
}

pub fn format_module_git<'a>(c: &mut Config,
                             last_successful: Option<&'a str>)
                             -> (Option<&'a str>, Option<ANSIString<'static>>) {
    use git2::Repository;
    use std::env;

    let mut output = String::new();

    let repo = Repository::discover(env::current_dir().unwrap());
    if let Ok(repo) = repo {
        if let Ok(reference) = repo.head() {
            // Current branch name
            output.push_str(reference.shorthand().unwrap());

            let show_diffs = c.get_bool("modules.git.show_diff_stats").unwrap_or_default();
            // Insertions and deletions
            if show_diffs {
                if let Ok(diff) = repo.diff_index_to_workdir(None, None) {
                    if let Ok(stats) = diff.stats() {
                        if stats.files_changed() > 0 {
                            let symbol_deletion = c.get_str("modules.git.symbol_deletion")
                                .unwrap_or_default();
                            let symbol_insertion = c.get_str("modules.git.symbol_insertion")
                                .unwrap_or_default();

                            output.push_str(&format!(" ({}{}, {}{})",
                                                     symbol_deletion,
                                                     stats.deletions(),
                                                     symbol_insertion,
                                                     stats.insertions()));
                        }
                    }
                }
            }
        }
    }

    if output.is_empty() {
        return (None, None);
    }

    // let symbol_clean = c.get_str("modules.git.symbol_clean").unwrap_or_default();
    // let symbol_push = c.get_str("modules.git.symbol_push").unwrap_or_default();
    // let symbol_pull = c.get_str("modules.git.symbol_pull").unwrap_or_default();

    format_module(c, "git", Some(output), last_successful)
}

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

    #[test]
    fn test_string_to_colour() {
        const CONTENT: &'static str = "hi";

        struct Test<'a> {
            input: &'static str,
            expected: &'a str,
        }

        let tests = [Test {
                         input: "green",
                         expected: &format!("\x1B[38;5;2m{}\x1B[0m", CONTENT),
                     },
                     Test {
                         input: "bright_green",
                         expected: &format!("\x1B[38;5;10m{}\x1B[0m", CONTENT),
                     }];

        for test in &tests {
            let result = string_to_colour(test.input.to_string());
            let result = format!("{}", result.paint(CONTENT));

            assert_eq!(test.expected, result);
        }
    }

    #[test]
    #[should_panic]
    fn test_string_to_colour_invalid_input() {
        struct Test {
            input: &'static str,
        }

        let tests = [Test { input: "green" }, Test { input: "invalid" }];

        for test in &tests {
            string_to_colour(test.input.to_string());
        }
    }
}
<|endoftext|>


<commit_before>use std::fmt;
use std::mem;
use std::ptr;
use std::slice;
use std::iter;
use std::io::{self, Write};
use std::cmp;
use std::rt::heap;
use std::cell::Cell;
use std::marker::PhantomData;
use std::intrinsics;

// A chunk in the vector
struct Chunk<T> {
    // Previous chunk
    prev: *mut Chunk<T>,
    // Next chunk
    next: *mut Chunk<T>,
    // Count of items
    len: usize,
    // Capacity
    cap: usize,
    // Items follow in memory
    items: [T; 0]
}

pub struct MonoVec<T> {
    head: Cell<*mut Chunk<T>>,
    tail: Cell<*mut Chunk<T>>,
    _ph: PhantomData<T>
}

unsafe impl<T: Send> Send for MonoVec<T> {}

impl<T> Chunk<T> {
    fn array_size(len: usize) -> usize {
        len.checked_mul(mem::size_of::<T>()).unwrap()
    }

    fn mem_size(len: usize) -> usize {
        mem::size_of::<Self>().checked_add(Self::array_size(len)).unwrap()
    }

    fn new(cap: usize) -> *mut Self {
        unsafe {
            let res = heap::allocate(Self::mem_size(cap),
                                     mem::align_of::<Self>()) as *mut Self;
            if res.is_null() {
                panic!("MonoVec: failed to allocate chunk!")
            }
            ptr::write(&mut (*res).prev, ptr::null_mut());
            ptr::write(&mut (*res).next, ptr::null_mut());
            ptr::write(&mut (*res).len, 0);
            ptr::write(&mut (*res).cap, cap);
            res
        }
    }
}

impl<T> MonoVec<T> {
    pub fn new() -> Self {
        Self::with_capacity(8)
    }

    pub fn with_capacity(cap: usize) -> Self {
        let head = Chunk::new(cmp::max(cap, 1));
        MonoVec {
            head: Cell::new(head),
            tail: Cell::new(head),
            _ph: PhantomData
        }
    }

    // FIXME: track total len in header to make this O(1)?
    pub fn len(&self) -> usize {
        let mut len = 0;
        let mut cur = self.head.get();

        while !cur.is_null() {
            unsafe {
                len += (*cur).len;
                cur = (*cur).next;
            }
        }
        len
    }

    // Reserves space for at least `len` more contiguous elements, returning
    // a pointer to the space and the available capacity (which may be > `len`)
    #[inline(never)]
    pub fn reserve(&self, len: usize) -> (*mut T, usize) {
        unsafe {
            let tail = self.tail.get();
            let cap = (*tail).cap;
            if cap - (*tail).len < len {
                // Grow capacity exponentially to amortize cost of insertions
                let mut new_cap = cap.checked_mul(2).unwrap();
                while new_cap < len {
                    new_cap = new_cap.checked_mul(2).unwrap();
                }
                let new = Chunk::new(new_cap);

                (*new).prev = tail;
                (*tail).next = new;
                self.tail.set(new);
            }

            let tail = self.tail.get();
            let ptr = (*tail).items.as_mut_ptr().offset((*tail).len as isize);
            let cap = (*tail).cap - (*tail).len;
            (ptr, cap)
        }
    }

    // Adds to length of curent chunk.  Usually used after
    // writing into reserved space.
    pub unsafe fn add_len(&self, len: usize) {
        let tail = self.tail.get();
        (*tail).len += len;
    }

    // Shrinks length of allocation at (ptr, ptr + old_len) if possible
    pub unsafe fn shrink_len(&self, ptr: *mut T, old_len: usize, new_len: usize) {
        let tail = self.tail.get();
        if ptr.offset(old_len as isize) == (*tail).items.as_mut_ptr().offset((*tail).len as isize) {
            (*tail).len = (*tail).len - old_len + new_len;
        }
    }

    #[inline]
    pub fn push(&self, elem: T) -> &T {
        let (ptr, _) = self.reserve(1);
        unsafe {
            ptr::write(ptr, elem);
            self.add_len(1);
            &*ptr
        }
    }

    pub fn extend_as_slice<E: IntoIterator<Item=T>>(&self, elems: E) -> &[T]
            where E::IntoIter: ExactSizeIterator {
        let iter = elems.into_iter();
        let len = iter.len();
        let (ptr, _) = self.reserve(len);
        let mut cur = ptr;
        unsafe {
            for elem in iter {
                ptr::write(cur, elem);
                cur = cur.offset(1);
            }
            self.add_len(len);
            slice::from_raw_parts(ptr, len)
        }
    }

    pub fn chunks(&self) -> Chunks<T> {
        Chunks {
            start: self.head.get(),
            end: self.tail.get(),
            _ph: PhantomData
        }
    }

    pub fn items(&self) -> Items<T> {
        #[inline(always)]
        fn id<T>(x: T) -> T { x }
        Items(self.chunks().flat_map(id))
    }
}

impl<'a, T: 'a> IntoIterator for &'a MonoVec<T> {
    type Item = &'a T;
    type IntoIter = Items<'a, T>;

    fn into_iter(self) -> Self::IntoIter {
        self.items()
    }
}

impl io::Write for MonoVec<u8> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let (ptr, len) = self.reserve(0);
        let len = cmp::min(len, buf.len());
        if len != 0 {
            unsafe {
                ptr::copy_nonoverlapping(buf.as_ptr(), ptr, len);
            }
        }
        Ok(len)
    }

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

impl<T> Drop for MonoVec<T> {
    fn drop(&mut self) {
        unsafe {
            let mut chunk = self.head.get();
            let mut next;
            while !chunk.is_null() {
                next = (*chunk).next;
                if intrinsics::needs_drop::<T>() {
                    let mut cur = (*chunk).items.as_mut_ptr();
                    let end = cur.offset((*chunk).len as isize);
                    while cur < end {
                        intrinsics::drop_in_place(cur);
                        cur = cur.offset(1);
                    }
                }
                heap::deallocate(chunk as *mut u8,
                                 (*chunk).cap * mem::size_of::<T>(),
                                 mem::min_align_of::<T>());
                chunk = next;
            }
        }
    }
}

pub struct Chunks<'a, T: 'a> {
    start: *mut Chunk<T>,
    end: *mut Chunk<T>,
    _ph: PhantomData<&'a [T]>
}

impl<'a, T> Iterator for Chunks<'a, T> {
    type Item = &'a [T];

    fn next(&mut self) -> Option<&'a [T]> {
        let chunk = self.start;
        if chunk.is_null() {
            None
        } else {
            unsafe {
                if chunk == self.end {
                    self.start = ptr::null_mut();
                    self.end = ptr::null_mut()
                } else {
                    self.start = (*chunk).next
                }
                Some(slice::from_raw_parts((*chunk).items.as_ptr(), (*chunk).len))
            }
        }
    }
}

unsafe impl <'a, T:Send> Send for Chunks<'a, T> {}
// FIXME: need to double check that this holds.
// It certainly doesn't for the vec itself.
unsafe impl <'a, T:Sync> Sync for Chunks<'a, T> {}

impl<'a, T> DoubleEndedIterator for Chunks<'a, T> {
    fn next_back(&mut self) -> Option<&'a [T]> {
        let chunk = self.end;
        if chunk.is_null() {
            None
        } else {
            unsafe {
                if chunk == self.start {
                    self.start = ptr::null_mut();
                    self.end = ptr::null_mut()
                } else {
                    self.end = (*chunk).prev
                }
                Some(slice::from_raw_parts((*chunk).items.as_ptr(), (*chunk).len))
            }
        }
    }
}

// Wrapper to hide ugly adapter type
pub struct Items<'a, T: 'a>(iter::FlatMap<Chunks<'a, T>, &'a [T], fn(&'a [T]) -> &'a [T]>);

impl<'a, T: 'a> Iterator for Items<'a, T> {
    type Item = &'a T;
    fn next(&mut self) -> Option<&'a T> {
        self.0.next()
    }
}

impl<'a, T: 'a> DoubleEndedIterator for Items<'a, T> {
    fn next_back(&mut self) -> Option<&'a T> {
        self.0.next_back()
    }
}

impl<T: fmt::Debug> fmt::Debug for MonoVec<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut need_comma = false;
        try!(write!(f, "["));
        for elem in self {
            if need_comma {
                try!(write!(f, ", "));
            }
            try!(elem.fmt(f));
            need_comma = true;
        }
        write!(f, "]")
    }
}
<commit_msg>Add clear method to MonoVec<commit_after>use std::fmt;
use std::mem;
use std::ptr;
use std::slice;
use std::iter;
use std::io::{self, Write};
use std::cmp;
use std::rt::heap;
use std::cell::Cell;
use std::marker::PhantomData;
use std::intrinsics;

// A chunk in the vector
struct Chunk<T> {
    // Previous chunk
    prev: *mut Chunk<T>,
    // Next chunk
    next: *mut Chunk<T>,
    // Count of items
    len: usize,
    // Capacity
    cap: usize,
    // Items follow in memory
    items: [T; 0]
}

pub struct MonoVec<T> {
    head: Cell<*mut Chunk<T>>,
    tail: Cell<*mut Chunk<T>>,
    _ph: PhantomData<T>
}

unsafe impl<T: Send> Send for MonoVec<T> {}

impl<T> Chunk<T> {
    fn array_size(len: usize) -> usize {
        len.checked_mul(mem::size_of::<T>()).unwrap()
    }

    fn mem_size(len: usize) -> usize {
        mem::size_of::<Self>().checked_add(Self::array_size(len)).unwrap()
    }

    fn new(cap: usize) -> *mut Self {
        unsafe {
            let res = heap::allocate(Self::mem_size(cap),
                                     mem::align_of::<Self>()) as *mut Self;
            if res.is_null() {
                panic!("MonoVec: failed to allocate chunk!")
            }
            ptr::write(&mut (*res).prev, ptr::null_mut());
            ptr::write(&mut (*res).next, ptr::null_mut());
            ptr::write(&mut (*res).len, 0);
            ptr::write(&mut (*res).cap, cap);
            res
        }
    }
}

impl<T> MonoVec<T> {
    pub fn new() -> Self {
        Self::with_capacity(8)
    }

    pub fn with_capacity(cap: usize) -> Self {
        let head = Chunk::new(cmp::max(cap, 1));
        MonoVec {
            head: Cell::new(head),
            tail: Cell::new(head),
            _ph: PhantomData
        }
    }

    // FIXME: track total len in header to make this O(1)?
    pub fn len(&self) -> usize {
        let mut len = 0;
        let mut cur = self.head.get();

        while !cur.is_null() {
            unsafe {
                len += (*cur).len;
                cur = (*cur).next;
            }
        }
        len
    }

    // Reserves space for at least `len` more contiguous elements, returning
    // a pointer to the space and the available capacity (which may be > `len`)
    #[inline(never)]
    pub fn reserve(&self, len: usize) -> (*mut T, usize) {
        unsafe {
            let tail = self.tail.get();
            let cap = (*tail).cap;
            if cap - (*tail).len < len {
                // Grow capacity exponentially to amortize cost of insertions
                let mut new_cap = cap.checked_mul(2).unwrap();
                while new_cap < len {
                    new_cap = new_cap.checked_mul(2).unwrap();
                }
                let new = Chunk::new(new_cap);

                (*new).prev = tail;
                (*tail).next = new;
                self.tail.set(new);
            }

            let tail = self.tail.get();
            let ptr = (*tail).items.as_mut_ptr().offset((*tail).len as isize);
            let cap = (*tail).cap - (*tail).len;
            (ptr, cap)
        }
    }

    // Adds to length of curent chunk.  Usually used after
    // writing into reserved space.
    pub unsafe fn add_len(&self, len: usize) {
        let tail = self.tail.get();
        (*tail).len += len;
    }

    // Shrinks length of allocation at (ptr, ptr + old_len) if possible
    pub unsafe fn shrink_len(&self, ptr: *mut T, old_len: usize, new_len: usize) {
        let tail = self.tail.get();
        if ptr.offset(old_len as isize) == (*tail).items.as_mut_ptr().offset((*tail).len as isize) {
            (*tail).len = (*tail).len - old_len + new_len;
        }
    }

    #[inline]
    pub fn push(&self, elem: T) -> &T {
        let (ptr, _) = self.reserve(1);
        unsafe {
            ptr::write(ptr, elem);
            self.add_len(1);
            &*ptr
        }
    }

    pub fn extend_as_slice<E: IntoIterator<Item=T>>(&self, elems: E) -> &[T]
            where E::IntoIter: ExactSizeIterator {
        let iter = elems.into_iter();
        let len = iter.len();
        let (ptr, _) = self.reserve(len);
        let mut cur = ptr;
        unsafe {
            for elem in iter {
                ptr::write(cur, elem);
                cur = cur.offset(1);
            }
            self.add_len(len);
            slice::from_raw_parts(ptr, len)
        }
    }

    pub fn clear(&mut self) {
        unsafe {
            loop {
                let chunk = self.head.get();
                self.head.set((*chunk).next);
                if intrinsics::needs_drop::<T>() {
                    let ptr = (*chunk).items.as_mut_ptr();
                    let end = ptr.offset((*chunk).len as isize);
                    while ptr < end {
                        intrinsics::drop_in_place(ptr);
                    }
                }
                if chunk == self.tail.get() {
                    break
                }
                heap::deallocate(chunk as *mut u8,
                                 mem::size_of::<Chunk<T>>() + (*chunk).len * mem::size_of::<T>(),
                                 mem::align_of::<Chunk<T>>());
            }
            let save = self.tail.get();
            self.head.set(save);
            (*save).len = 0;
        }
    }

    pub fn chunks(&self) -> Chunks<T> {
        Chunks {
            start: self.head.get(),
            end: self.tail.get(),
            _ph: PhantomData
        }
    }

    pub fn items(&self) -> Items<T> {
        #[inline(always)]
        fn id<T>(x: T) -> T { x }
        Items(self.chunks().flat_map(id))
    }
}

impl<'a, T: 'a> IntoIterator for &'a MonoVec<T> {
    type Item = &'a T;
    type IntoIter = Items<'a, T>;

    fn into_iter(self) -> Self::IntoIter {
        self.items()
    }
}

impl io::Write for MonoVec<u8> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let (ptr, len) = self.reserve(0);
        let len = cmp::min(len, buf.len());
        if len != 0 {
            unsafe {
                ptr::copy_nonoverlapping(buf.as_ptr(), ptr, len);
            }
        }
        Ok(len)
    }

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

impl<T> Drop for MonoVec<T> {
    fn drop(&mut self) {
        self.clear();
        let chunk = self.head.get();
        unsafe {
            heap::deallocate(chunk as *mut u8,
                             mem::size_of::<Chunk<T>>() + (*chunk).len * mem::size_of::<T>(),
                             mem::align_of::<Chunk<T>>());

        }
    }
}

pub struct Chunks<'a, T: 'a> {
    start: *mut Chunk<T>,
    end: *mut Chunk<T>,
    _ph: PhantomData<&'a [T]>
}

impl<'a, T> Iterator for Chunks<'a, T> {
    type Item = &'a [T];

    fn next(&mut self) -> Option<&'a [T]> {
        let chunk = self.start;
        if chunk.is_null() {
            None
        } else {
            unsafe {
                if chunk == self.end {
                    self.start = ptr::null_mut();
                    self.end = ptr::null_mut()
                } else {
                    self.start = (*chunk).next
                }
                Some(slice::from_raw_parts((*chunk).items.as_ptr(), (*chunk).len))
            }
        }
    }
}

unsafe impl <'a, T:Send> Send for Chunks<'a, T> {}
// FIXME: need to double check that this holds.
// It certainly doesn't for the vec itself.
unsafe impl <'a, T:Sync> Sync for Chunks<'a, T> {}

impl<'a, T> DoubleEndedIterator for Chunks<'a, T> {
    fn next_back(&mut self) -> Option<&'a [T]> {
        let chunk = self.end;
        if chunk.is_null() {
            None
        } else {
            unsafe {
                if chunk == self.start {
                    self.start = ptr::null_mut();
                    self.end = ptr::null_mut()
                } else {
                    self.end = (*chunk).prev
                }
                Some(slice::from_raw_parts((*chunk).items.as_ptr(), (*chunk).len))
            }
        }
    }
}

// Wrapper to hide ugly adapter type
pub struct Items<'a, T: 'a>(iter::FlatMap<Chunks<'a, T>, &'a [T], fn(&'a [T]) -> &'a [T]>);

impl<'a, T: 'a> Iterator for Items<'a, T> {
    type Item = &'a T;
    fn next(&mut self) -> Option<&'a T> {
        self.0.next()
    }
}

impl<'a, T: 'a> DoubleEndedIterator for Items<'a, T> {
    fn next_back(&mut self) -> Option<&'a T> {
        self.0.next_back()
    }
}

impl<T: fmt::Debug> fmt::Debug for MonoVec<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut need_comma = false;
        try!(write!(f, "["));
        for elem in self {
            if need_comma {
                try!(write!(f, ", "));
            }
            try!(elem.fmt(f));
            need_comma = true;
        }
        write!(f, "]")
    }
}
<|endoftext|>