<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;
extern crate chrono;

extern crate mimirsbrunn;
extern crate serde;
extern crate geo;

#[macro_use]
extern crate mdo;

use std::collections::{BTreeSet, BTreeMap};
use mimir::rubber::Rubber;

pub type AdminsVec = Vec<Rc<mimir::Admin>>;
pub type StreetsVec = Vec<mimir::Street>;
pub type OsmPbfReader = osmpbfreader::OsmPbfReader<std::fs::File>;
pub type StreetWithRelationSet = BTreeSet<osmpbfreader::OsmId>;
pub type AdminSet = BTreeSet<Rc<mimir::Admin>>;
pub type NameAdminMap = BTreeMap<StreetKey, Vec<osmpbfreader::OsmId>>;

use mimirsbrunn::admin_geofinder::AdminGeoFinder;
use std::rc::Rc;
use std::cell::Cell;

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

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

Options:
    -h, --help              Show this message.
    -i, --input=<file>      OSM PBF file.
    -l, --level=<level>     Admin levels to keep.
    -C, --city-level=<level>
                            City level to calculate weight, [default: 8]
    -w, --import-way        Import ways
    -a, --import-admin      Import admins
    -c, --connection-string=<connection-string>
                            Elasticsearch parameters, [default: http://localhost:9200/munin]
    -d, --dataset=<dataset>
                            Name of the dataset, [default: fr]

";

#[derive(Debug)]
struct AdminMatcher {
    admin_levels: BTreeSet<u32>,
}
impl AdminMatcher {
    pub fn new(levels: BTreeSet<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,
        }
    }
}

#[derive(Debug, Clone, Ord, PartialOrd, Eq, PartialEq)]
pub struct StreetKey {
    pub name: String,
    pub admins: AdminSet,
}

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

fn administrative_regions(pbf: &mut OsmPbfReader, levels: BTreeSet<u32>) -> AdminsVec {
    let mut administrative_regions = AdminsVec::new();
    let matcher = AdminMatcher::new(levels);
    info!("reading pbf...");
    let objects = osmpbfreader::get_objs_and_deps(pbf, |o| matcher.is_admin(o)).unwrap();
    info!("reading pbf done.");
    // 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::CoordWrapper(geo::Coordinate {
                                                           x: node.lat,
                                                           y: node.lon,
                                                       }))
                                                   }
                                                   _ => None,
                                               }
                                           })
                                       });
            let (admin_id, insee_id) = match relation.tags.get("ref:INSEE") {
                Some(val) => {
                    (format!("admin:fr:{}", val.trim_left_matches('0')), 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 = mimirsbrunn::boundaries::build_boundary(&relation, &objects);
            let admin = mimir::Admin {
                id: admin_id,
                insee: insee_id.to_string(),
                level: level,
                label: name.to_string(),
                zip_code: zip_code.to_string(),
                // TODO weight value ?
                weight: Cell::new(0),
                coord: coord_centre,
                boundary: boundary,
            };
            administrative_regions.push(Rc::new(admin));
        }
    }
    return administrative_regions;
}

fn make_admin_geofinder(admins: &AdminsVec) -> AdminGeoFinder {
    let mut geofinder = AdminGeoFinder::new();

    for a in admins {
        geofinder.insert(a.clone());
    }
    geofinder
}

fn get_street_admin(admins_geofinder: &AdminGeoFinder,
                    obj_map: &BTreeMap<osmpbfreader::OsmId, osmpbfreader::OsmObj>,
                    way: &osmpbfreader::objects::Way)
                    -> Vec<Rc<mimir::Admin>> {
    // for the moment we consider that the coord of the way is the coord of it's first node
    let coord = way.nodes
                   .iter()
                   .filter_map(|node_id| obj_map.get(&osmpbfreader::OsmId::Node(*node_id)))
                   .filter_map(|node_obj| {
                       if let &osmpbfreader::OsmObj::Node(ref node) = node_obj {
                           Some(geo::Coordinate {
                               x: node.lat,
                               y: node.lon,
                           })
                       } else {
                           None
                       }
                   })
                   .next();
    coord.map_or(vec![], |c| admins_geofinder.get(&c))
}

fn streets(pbf: &mut OsmPbfReader, admins: &AdminsVec, city_level: u32) -> StreetsVec {
    let admins_geofinder = make_admin_geofinder(admins);

    let is_valid_obj = |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())
            }
            osmpbfreader::OsmObj::Relation(ref rel) => {
                rel.tags.get("type").map_or(false, |v| v == "associatedStreet")
            }
            _ => false,
        }
    };
    info!("reading pbf...");
    let objs_map = osmpbfreader::get_objs_and_deps(pbf, is_valid_obj).unwrap();
    info!("reading pbf done.");
    let mut street_rel: StreetWithRelationSet = BTreeSet::new();
    let mut street_list: StreetsVec = vec![];
    // 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 duplicates in autocompleion,
    // we should tag the osm ways in the relation not to index them twice.

    for (_, rel_obj) in &objs_map {
        if let &osmpbfreader::OsmObj::Relation(ref rel) = rel_obj {
            let way_name = rel.tags.get("name");
            for ref_obj in &rel.refs {

                use mdo::option::*;
                let objs_map = &objs_map;
                let street_list = &mut street_list;
                let admins_geofinder = &admins_geofinder;

                let inserted = mdo! {
                    when ref_obj.member.is_way();
                    when ref_obj.role == "street";
                    obj =<< objs_map.get(&ref_obj.member);
                    way =<<obj.way();
                    way_name =<< way_name.or_else(|| way.tags.get("name"));
                    let admin = get_street_admin(admins_geofinder, objs_map, way);
                    ret ret(street_list.push(mimir::Street {
                                id: way.id.to_string(),
                                street_name: way_name.to_string(),
                                label: way_name.to_string(),
                                weight: 1,
                                administrative_regions: admin,
                    }))
                };
                if inserted.is_some() {
                    break;
                }
            }

            // Add osmid of all the relation members in de set
            // We don't create any street for all the osmid present in street_rel
            for ref_obj in &rel.refs {
                if ref_obj.member.is_way() {
                    street_rel.insert(ref_obj.member);
                }
            }
        };
    }

    // we merge all the ways with a key = way_name + admin list of level(=city_level)
    // we use a map NameAdminMap <key, value> to manage the merging of ways
    let mut name_admin_map: NameAdminMap = BTreeMap::new();
    for (osmid, obj) in &objs_map {
        if street_rel.contains(osmid) {
            continue;
        }
        use mdo::option::*;
        let admins_geofinder = &admins_geofinder;
        let objs_map = &objs_map;
        let name_admin_map = &mut name_admin_map;
        mdo! {
            way =<< obj.way();
            let admins: BTreeSet<Rc<mimir::Admin>> = get_street_admin(admins_geofinder, objs_map, way)
            	.into_iter()
            	.filter(|admin| admin.level == city_level)
            	.collect();

            way_name =<< way.tags.get("name");
            let key = StreetKey{name: way_name.to_string(), admins: admins};
            ret ret(name_admin_map.entry(key).or_insert(vec![]).push(*osmid))
        };
    }

    // Create a street for each way with osmid present in in objs_map
    for (_, way_ids) in name_admin_map {
        use mdo::option::*;
        let objs_map = &objs_map;
        let street_list = &mut street_list;
        let admins_geofinder = &admins_geofinder;
        mdo! {
            obj =<< objs_map.get(&way_ids[0]);
            way =<< obj.way();
            way_name =<< way.tags.get("name");
            let admins = get_street_admin(admins_geofinder, objs_map, way);
            ret ret(street_list.push(mimir::Street{
   	                    id: way.id.to_string(),
   	                    street_name: way_name.to_string(),
   	                    label: way_name.to_string(),
   	                    weight: 1,
   	                    administrative_regions: admins,
            }))
   	    };
    }

    street_list
}

fn main() {
    mimir::logger_init().unwrap();
    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 city_level = args.flag_city_level;
    let mut parsed_pbf = parse_osm_pbf(&args.flag_input);
    debug!("creation of indexes");
    let mut rubber = Rubber::new(&args.flag_connection_string);

    info!("creating adminstrative regions");
    let admins = administrative_regions(&mut parsed_pbf, levels);

    info!("computing city weight");
    let mut streets = streets(&mut parsed_pbf, &admins, city_level);

    for st in &mut streets {
        for admin in &mut st.administrative_regions {
            if admin.level == city_level {
                admin.weight.set(admin.weight.get() + 1)
            }
        }
    }

    for st in &mut streets {
        for admin in &mut st.administrative_regions {
            if admin.level == city_level {
                st.weight = admin.weight.get();
                break;
            }
        }
    }

    if args.flag_import_way {
        info!("importing streets into Mimir");
        let nb_streets = rubber.index("way", &args.flag_dataset, streets.into_iter())
                               .unwrap();
        info!("Nb of indexed street: {}", nb_streets);
    }

    let nb_admins = rubber.index("admin", &args.flag_dataset, admins.iter())
                          .unwrap();
    info!("Nb of indexed admin: {}", nb_admins);

}
<commit_msg>mimir: tabulation<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;
extern crate chrono;

extern crate mimirsbrunn;
extern crate serde;
extern crate geo;

#[macro_use]
extern crate mdo;

use std::collections::{BTreeSet, BTreeMap};
use mimir::rubber::Rubber;

pub type AdminsVec = Vec<Rc<mimir::Admin>>;
pub type StreetsVec = Vec<mimir::Street>;
pub type OsmPbfReader = osmpbfreader::OsmPbfReader<std::fs::File>;
pub type StreetWithRelationSet = BTreeSet<osmpbfreader::OsmId>;
pub type AdminSet = BTreeSet<Rc<mimir::Admin>>;
pub type NameAdminMap = BTreeMap<StreetKey, Vec<osmpbfreader::OsmId>>;

use mimirsbrunn::admin_geofinder::AdminGeoFinder;
use std::rc::Rc;
use std::cell::Cell;

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

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

Options:
    -h, --help              Show this message.
    -i, --input=<file>      OSM PBF file.
    -l, --level=<level>     Admin levels to keep.
    -C, --city-level=<level>
                            City level to calculate weight, [default: 8]
    -w, --import-way        Import ways
    -a, --import-admin      Import admins
    -c, --connection-string=<connection-string>
                            Elasticsearch parameters, [default: http://localhost:9200/munin]
    -d, --dataset=<dataset>
                            Name of the dataset, [default: fr]

";

#[derive(Debug)]
struct AdminMatcher {
    admin_levels: BTreeSet<u32>,
}
impl AdminMatcher {
    pub fn new(levels: BTreeSet<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,
        }
    }
}

#[derive(Debug, Clone, Ord, PartialOrd, Eq, PartialEq)]
pub struct StreetKey {
    pub name: String,
    pub admins: AdminSet,
}

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

fn administrative_regions(pbf: &mut OsmPbfReader, levels: BTreeSet<u32>) -> AdminsVec {
    let mut administrative_regions = AdminsVec::new();
    let matcher = AdminMatcher::new(levels);
    info!("reading pbf...");
    let objects = osmpbfreader::get_objs_and_deps(pbf, |o| matcher.is_admin(o)).unwrap();
    info!("reading pbf done.");
    // 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::CoordWrapper(geo::Coordinate {
                                                           x: node.lat,
                                                           y: node.lon,
                                                       }))
                                                   }
                                                   _ => None,
                                               }
                                           })
                                       });
            let (admin_id, insee_id) = match relation.tags.get("ref:INSEE") {
                Some(val) => {
                    (format!("admin:fr:{}", val.trim_left_matches('0')), 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 = mimirsbrunn::boundaries::build_boundary(&relation, &objects);
            let admin = mimir::Admin {
                id: admin_id,
                insee: insee_id.to_string(),
                level: level,
                label: name.to_string(),
                zip_code: zip_code.to_string(),
                // TODO weight value ?
                weight: Cell::new(0),
                coord: coord_centre,
                boundary: boundary,
            };
            administrative_regions.push(Rc::new(admin));
        }
    }
    return administrative_regions;
}

fn make_admin_geofinder(admins: &AdminsVec) -> AdminGeoFinder {
    let mut geofinder = AdminGeoFinder::new();

    for a in admins {
        geofinder.insert(a.clone());
    }
    geofinder
}

fn get_street_admin(admins_geofinder: &AdminGeoFinder,
                    obj_map: &BTreeMap<osmpbfreader::OsmId, osmpbfreader::OsmObj>,
                    way: &osmpbfreader::objects::Way)
                    -> Vec<Rc<mimir::Admin>> {
    // for the moment we consider that the coord of the way is the coord of it's first node
    let coord = way.nodes
                   .iter()
                   .filter_map(|node_id| obj_map.get(&osmpbfreader::OsmId::Node(*node_id)))
                   .filter_map(|node_obj| {
                       if let &osmpbfreader::OsmObj::Node(ref node) = node_obj {
                           Some(geo::Coordinate {
                               x: node.lat,
                               y: node.lon,
                           })
                       } else {
                           None
                       }
                   })
                   .next();
    coord.map_or(vec![], |c| admins_geofinder.get(&c))
}

fn streets(pbf: &mut OsmPbfReader, admins: &AdminsVec, city_level: u32) -> StreetsVec {
    let admins_geofinder = make_admin_geofinder(admins);

    let is_valid_obj = |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())
            }
            osmpbfreader::OsmObj::Relation(ref rel) => {
                rel.tags.get("type").map_or(false, |v| v == "associatedStreet")
            }
            _ => false,
        }
    };
    info!("reading pbf...");
    let objs_map = osmpbfreader::get_objs_and_deps(pbf, is_valid_obj).unwrap();
    info!("reading pbf done.");
    let mut street_rel: StreetWithRelationSet = BTreeSet::new();
    let mut street_list: StreetsVec = vec![];
    // 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 duplicates in autocompleion,
    // we should tag the osm ways in the relation not to index them twice.

    for (_, rel_obj) in &objs_map {
        if let &osmpbfreader::OsmObj::Relation(ref rel) = rel_obj {
            let way_name = rel.tags.get("name");
            for ref_obj in &rel.refs {

                use mdo::option::*;
                let objs_map = &objs_map;
                let street_list = &mut street_list;
                let admins_geofinder = &admins_geofinder;

                let inserted = mdo! {
                    when ref_obj.member.is_way();
                    when ref_obj.role == "street";
                    obj =<< objs_map.get(&ref_obj.member);
                    way =<<obj.way();
                    way_name =<< way_name.or_else(|| way.tags.get("name"));
                    let admin = get_street_admin(admins_geofinder, objs_map, way);
                    ret ret(street_list.push(mimir::Street {
                                id: way.id.to_string(),
                                street_name: way_name.to_string(),
                                label: way_name.to_string(),
                                weight: 1,
                                administrative_regions: admin,
                    }))
                };
                if inserted.is_some() {
                    break;
                }
            }

            // Add osmid of all the relation members in de set
            // We don't create any street for all the osmid present in street_rel
            for ref_obj in &rel.refs {
                if ref_obj.member.is_way() {
                    street_rel.insert(ref_obj.member);
                }
            }
        };
    }

    // we merge all the ways with a key = way_name + admin list of level(=city_level)
    // we use a map NameAdminMap <key, value> to manage the merging of ways
    let mut name_admin_map: NameAdminMap = BTreeMap::new();
    for (osmid, obj) in &objs_map {
        if street_rel.contains(osmid) {
            continue;
        }
        use mdo::option::*;
        let admins_geofinder = &admins_geofinder;
        let objs_map = &objs_map;
        let name_admin_map = &mut name_admin_map;
        mdo! {
            way =<< obj.way();
            let admins: BTreeSet<Rc<mimir::Admin>> = get_street_admin(admins_geofinder, objs_map, way)
            	.into_iter()
            	.filter(|admin| admin.level == city_level)
            	.collect();

            way_name =<< way.tags.get("name");
            let key = StreetKey{name: way_name.to_string(), admins: admins};
            ret ret(name_admin_map.entry(key).or_insert(vec![]).push(*osmid))
        };
    }

    // Create a street for each way with osmid present in in objs_map
    for (_, way_ids) in name_admin_map {
        use mdo::option::*;
        let objs_map = &objs_map;
        let street_list = &mut street_list;
        let admins_geofinder = &admins_geofinder;
        mdo! {
            obj =<< objs_map.get(&way_ids[0]);
            way =<< obj.way();
            way_name =<< way.tags.get("name");
            let admins = get_street_admin(admins_geofinder, objs_map, way);
            ret ret(street_list.push(mimir::Street{
   	                    id: way.id.to_string(),
   	                    street_name: way_name.to_string(),
   	                    label: way_name.to_string(),
   	                    weight: 1,
   	                    administrative_regions: admins,
            }))
        };
    }

    street_list
}

fn main() {
    mimir::logger_init().unwrap();
    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 city_level = args.flag_city_level;
    let mut parsed_pbf = parse_osm_pbf(&args.flag_input);
    debug!("creation of indexes");
    let mut rubber = Rubber::new(&args.flag_connection_string);

    info!("creating adminstrative regions");
    let admins = administrative_regions(&mut parsed_pbf, levels);

    info!("computing city weight");
    let mut streets = streets(&mut parsed_pbf, &admins, city_level);

    for st in &mut streets {
        for admin in &mut st.administrative_regions {
            if admin.level == city_level {
                admin.weight.set(admin.weight.get() + 1)
            }
        }
    }

    for st in &mut streets {
        for admin in &mut st.administrative_regions {
            if admin.level == city_level {
                st.weight = admin.weight.get();
                break;
            }
        }
    }

    if args.flag_import_way {
        info!("importing streets into Mimir");
        let nb_streets = rubber.index("way", &args.flag_dataset, streets.into_iter())
                               .unwrap();
        info!("Nb of indexed street: {}", nb_streets);
    }

    let nb_admins = rubber.index("admin", &args.flag_dataset, admins.iter())
                          .unwrap();
    info!("Nb of indexed admin: {}", nb_admins);

}
<|endoftext|>


<commit_before>extern crate libcratesio;

use libcratesio::{CratesIO, Crate, Error, ErrorKind};

#[test]
fn basic_data() {
    let krate = Crate::by_name("requests").unwrap();
    println!("{:?}", krate);
    assert_eq!(krate.name, "requests");
}

#[test]
fn raw_data() {
    let krate = CratesIO::query("requests").unwrap().as_json().unwrap();
    println!("{:#?}", krate);
    assert!(true);
}

#[test]
fn versions_iter() {
    let krate = Crate::by_name("requests").unwrap();
    for ver in krate.versions.iter() {
        println!("{:?}", ver);
    }
    assert!(true);
}

#[test]
fn error() {
    let no_such_crate = Crate::by_name("no_such_crate");
    assert!(no_such_crate.is_err());
    let err = no_such_crate.unwrap_err();
    if let Error(ErrorKind::CratesIOError(response), _) = err {
        assert_eq!(response.detail(), "Not Found");
    } else {
        assert!(false);
    }
}
<commit_msg>Update tests<commit_after>extern crate libcratesio;

use libcratesio::{CratesIO, Crate, Error, ErrorKind};

#[test]
fn basic_data() {
    let krate = Crate::by_name("libcratesio").unwrap();
    println!("{:?}", krate);
    assert_eq!(krate.name, "libcratesio");
}

#[test]
fn raw_data() {
    let krate = CratesIO::query("libcratesio")
        .unwrap()
        .as_json()
        .unwrap();
    println!("{:#?}", krate);
    assert!(true);
}

#[test]
fn versions_iter() {
    let krate = Crate::by_name("libcratesio").unwrap();
    for ver in krate.versions.iter() {
        println!("{:?}", ver);
    }
    assert!(true);
}

#[test]
fn error() {
    let no_such_crate = Crate::by_name("no_such_crate");
    assert!(no_such_crate.is_err());
    let err = no_such_crate.unwrap_err();
    if let Error(ErrorKind::CratesIOError(response), _) = err {
        assert_eq!(response.detail(), "Not Found");
    } else {
        assert!(false);
    }
}
<|endoftext|>


<commit_before>use std::collections::{BTreeMap, HashMap};
use std::path::{Path, PathBuf};

#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct BranchContext {
    files: HashMap<PathBuf, BranchAnalysis>,
}

impl BranchContext {
    pub fn is_branch<P: AsRef<Path>>(&self, path: P, line: usize) -> bool {
        if let Some(file) = self.files.get(path.as_ref()) {
            file.is_branch(line)
        } else {
            false
        }
    }
}

/// Coverage context for all the branches
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct BranchAnalysis {
    /// Each key is `LineRange` showing a region of the code containing a set of branches with the
    /// value being a `LineRange` for each branch in the code
    /// TODO consider BTreeMap then can order on line range start
    branches: BTreeMap<LineRange, Branches>,
}

impl BranchAnalysis {
    /// Returns true if the line is part of a branch
    pub fn is_branch(&self, line: usize) -> bool {
        self.branches.iter().any(|(k, _)| k.contains(line))
    }
}

/// Represents possible branches through an execution
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct Branches {
    /// Line ranges for each branch
    ranges: Vec<LineRange>,
    /// Whether there is an implicit or empty default branch i.e. missing or empty `else` in an
    /// `if` statement
    implicit_default: bool,
}

/// The start and end of contiguous range of lines. The range is contained within
/// `start..end`
#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash, Ord, PartialOrd)]
pub struct LineRange {
    /// Start of the line range (inclusive)
    start: usize,
    /// End of the line range (exclusive)
    end: usize,
}

impl LineRange {
    /// Returns true if the line is contained within the line range
    pub fn contains(&self, line: usize) -> bool {
        line >= self.start && line < self.end
    }
}
<commit_msg>Work out branches for loops<commit_after>use std::collections::{BTreeMap, HashMap};
use std::path::{Path, PathBuf};
use syn::*;

#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct BranchContext {
    files: HashMap<PathBuf, BranchAnalysis>,
}

impl BranchContext {
    pub fn is_branch<P: AsRef<Path>>(&self, path: P, line: usize) -> bool {
        if let Some(file) = self.files.get(path.as_ref()) {
            file.is_branch(line)
        } else {
            false
        }
    }
}

/// Coverage context for all the branches
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct BranchAnalysis {
    /// Each key is `LineRange` showing a region of the code containing a set of branches with the
    /// value being a `LineRange` for each branch in the code
    /// TODO consider BTreeMap then can order on line range start
    branches: BTreeMap<LineRange, Branches>,
}

impl BranchAnalysis {
    /// Returns true if the line is part of a branch
    pub fn is_branch(&self, line: usize) -> bool {
        self.branches.iter().any(|(k, _)| k.contains(line))
    }

    pub fn register_expr(&mut self, expr: &Expr) {
        let range = LineRange::from(expr);
        match *expr {
            Expr::If(ref e) => {
                self.branches.insert(range, e.into());
            }
            Expr::Match(ref e) => {
                self.branches.insert(range, e.into());
            }
            Expr::ForLoop(ref e) => {
                self.branches.insert(range, e.into());
            }
            Expr::While(ref e) => {
                self.branches.insert(range, e.into());
            }
            _ => {}
        }
    }
}

/// Represents possible branches through an execution
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct Branches {
    /// Line ranges for each branch
    ranges: Vec<LineRange>,
    /// Whether there is an implicit or empty default branch i.e. missing or empty `else` in an
    /// `if` statement
    implicit_default: bool,
}

impl From<&ExprIf> for Branches {
    fn from(expr: &ExprIf) -> Self {
        todo!()
    }
}

impl From<&ExprMatch> for Branches {
    fn from(expr: &ExprMatch) -> Self {
        todo!()
    }
}

impl From<&ExprForLoop> for Branches {
    fn from(expr: &ExprForLoop) -> Self {
        Self {
            implicit_default: true,
            ranges: vec![expr.body.clone().into()],
        }
    }
}

impl From<&ExprWhile> for Branches {
    fn from(expr: &ExprWhile) -> Self {
        Self {
            implicit_default: true,
            ranges: vec![expr.body.clone().into()],
        }
    }
}

/// The start and end of contiguous range of lines. The range is contained within
/// `start..end`
#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash, Ord, PartialOrd)]
pub struct LineRange {
    /// Start of the line range (inclusive)
    start: usize,
    /// End of the line range (exclusive)
    end: usize,
}

impl<T> From<T> for LineRange
where
    T: spanned::Spanned,
{
    fn from(t: T) -> Self {
        Self {
            start: t.span().start().line,
            end: t.span().end().line + 1,
        }
    }
}

impl LineRange {
    /// Returns true if the line is contained within the line range
    pub fn contains(&self, line: usize) -> bool {
        line >= self.start && line < self.end
    }
}
<|endoftext|>


<commit_before>//! Licensed under the Apache License, Version 2.0
//! http://www.apache.org/licenses/LICENSE-2.0 or the MIT license
//! http://opensource.org/licenses/MIT, at your
//! option. This file may not be copied, modified, or distributed
//! except according to those terms.

#![feature(core, test)]

#[macro_use]
extern crate itertools;

extern crate test;

use std::fmt::Debug;
use std::iter::order;
use itertools::Itertools;
use itertools::Interleave;
use itertools::Zip;

use itertools as it;

fn assert_iters_equal<A, I, J>(mut it: I, mut jt: J) where
    A: PartialEq + Debug,
    I: Iterator<Item=A>,
    J: Iterator<Item=A>,
{
    loop {
        let elti = it.next();
        let eltj = jt.next();
        assert_eq!(elti, eltj);
        if elti.is_none() { break; }
    }
}

#[test]
fn product2() {
    let s = "αβ";

    let mut prod = iproduct!(s.chars(), range(0, 2));
    assert!(prod.next() == Some(('α', 0)));
    assert!(prod.next() == Some(('α', 1)));
    assert!(prod.next() == Some(('β', 0)));
    assert!(prod.next() == Some(('β', 1)));
    assert!(prod.next() == None);
}

#[test]
fn product3() {
    let prod = iproduct!(range(0, 3), range(0, 2), range(0, 2));
    assert_eq!(prod.size_hint(), (12, Some(12)));
    let v = prod.collect_vec();
    for i in range(0,3) {
        for j in range(0, 2) {
            for k in range(0, 2) {
                assert!((i, j, k) == v[(i * 2 * 2 + j * 2 + k) as usize]);
            }
        }
    }
    for (_, _, _, _) in iproduct!(range(0, 3), range(0, 2), range(0, 2), range(0, 3)) {
        /* test compiles */
    }
}

#[test]
fn izip3() {
    let mut zip = Zip::new((range(0, 3), range(0, 2), range(0, 2i8)));
    for i in range(0, 2) {
        assert!((i as usize, i, i as i8) == zip.next().unwrap());
    }
    assert!(zip.next().is_none());

    let xs: [isize; 0] = [];
    let mut zip = Zip::new((range(0, 3), range(0, 2), range(0, 2i8), xs.iter()));
    assert!(zip.next().is_none());

    for (_, _, _, _) in Zip::new((range(0, 3), range(0, 2), range(0, 2), range(0, 3))) {
        /* test compiles */
    }
}

#[test]
fn write_to() {
    let xs = [7, 9, 8];
    let mut ys = [0; 5];
    let cnt = it::write(ys.iter_mut(), xs.iter().map(|x| *x));
    assert!(cnt == xs.len());
    assert!(ys == [7, 9, 8, 0, 0]);

    let cnt = it::write(ys.iter_mut(), range(0,10));
    assert!(cnt == ys.len());
    assert!(ys == [0, 1, 2, 3, 4]);
}

#[test]
fn interleave() {
    let xs: [u8; 0]  = [];
    let ys = [7u8, 9, 8, 10];
    let zs = [2u8, 77];
    let it = Interleave::new(xs.iter(), ys.iter());
    assert!(order::eq(it, ys.iter()));

    let rs = [7u8, 2, 9, 77, 8, 10];
    let it = Interleave::new(ys.iter(), zs.iter());
    assert!(order::eq(it, rs.iter()));
}

#[test]
fn times() {
    assert!(it::times(0).count() == 0);
    assert!(it::times(5).count() == 5);
}

#[test]
fn drain() {
    let xs = [1i32,2,3];
    let mut sum = 0;
    xs.iter().map(|elt| sum += *elt).drain();
    assert!(sum == 6);
}

#[test]
fn foreach() {
    let xs = [1i32, 2, 3];
    let mut sum = 0;
    xs.iter().foreach(|elt| sum += *elt);
    assert!(sum == 6);
}

#[test]
fn dropn() {
    let xs = [1, 2, 3];
    let mut it = xs.iter();
    assert!(it.dropn(2) == 2);
    assert!(it.next().is_some());
    assert!(it.next().is_none());
    let mut it = xs.iter();
    assert!(it.dropn(5) == 3);
    assert!(it.next().is_none());
}

#[test]
fn dropping() {
    let xs = [1, 2, 3];
    let mut it = xs.iter().dropping(2);
    assert!(it.next().is_some());
    assert!(it.next().is_none());
    let mut it = xs.iter().dropping(5);
    assert!(it.next().is_none());
}

#[test]
fn intersperse() {
    let xs = ["a", "", "b", "c"];
    let v: Vec<&str> = xs.iter().map(|x| x.clone()).intersperse(", ").collect();
    let text: String = v.concat();
    assert_eq!(text, "a, , b, c".to_string());

    let ys = [0, 1, 2, 3];
    let mut it = ys[..0].iter().map(|x| *x).intersperse(1);
    assert!(it.next() == None);
}

#[test]
fn linspace() {
    let mut iter = it::linspace::<f32>(0., 2., 3);
    assert_eq!(iter.next(), Some(0.0));
    assert_eq!(iter.next(), Some(1.0));
    assert_eq!(iter.next(), Some(2.0));
    assert_eq!(iter.next(), None);

    let mut iter = it::linspace::<f32>(0., -2., 4);
    assert_eq!(iter.next(), Some(0.));
    assert_eq!(iter.next(), Some(-0.666666666667));
    assert_eq!(iter.next(), Some(-1.333333333333));
    assert_eq!(iter.next(), Some(-2.));
    assert_eq!(iter.next(), None);

    let mut iter = it::linspace::<f32>(0., 1., 1);
    assert_eq!(iter.next(), Some(0.));
    assert_eq!(iter.next(), None);

    let mut iter = it::linspace::<f32>(0., 1., 0);
    assert_eq!(iter.next(), None);
}

#[test]
fn dedup() {
    let xs = [0, 1, 1, 1, 2, 1, 3, 3];
    let ys = [0, 1, 2, 1, 3];
    assert_iters_equal(ys.iter(), xs.iter().dedup());
    let xs = [0, 0, 0, 0, 0];
    let ys = [0];
    assert_iters_equal(ys.iter(), xs.iter().dedup());
}

#[test]
fn batching() {
    let xs = [0, 1, 2, 1, 3];
    let ys = [(0, 1), (2, 1)];

    // An iterator that gathers elements up in pairs
    let pit = xs.iter().cloned().batching(|mut it| {
               match it.next() {
                   None => None,
                   Some(x) => match it.next() {
                       None => None,
                       Some(y) => Some((x, y)),
                   }
               }
           });
    assert_iters_equal(pit, ys.iter().cloned());
}

#[test]
fn group_by() {
    let xs = [0, 1, 1, 1, 2, 1, 3, 3];
    let ans = vec![(0, vec![0]), (1, vec![1, 1, 1]),
                   (2, vec![2]), (1, vec![1]), (3, vec![3, 3])];

    let gb = xs.iter().cloned().group_by(|elt| *elt);
    assert_iters_equal(gb, ans.into_iter());
}

#[test]
fn put_back() {
    let xs = [0, 1, 1, 1, 2, 1, 3, 3];
    let mut pb = it::PutBack::new(xs.iter().cloned());
    pb.next();
    pb.put_back(1);
    pb.put_back(0);
    assert_iters_equal(pb, xs.iter().cloned());
}

#[test]
fn tee() {
    let xs  = [0, 1, 2, 3];
    let (mut t1, mut t2) = xs.iter().cloned().tee();
    assert_eq!(t1.next(), Some(0));
    assert_eq!(t2.next(), Some(0));
    assert_eq!(t1.next(), Some(1));
    assert_eq!(t1.next(), Some(2));
    assert_eq!(t1.next(), Some(3));
    assert_eq!(t1.next(), None);
    assert_eq!(t2.next(), Some(1));
    assert_eq!(t2.next(), Some(2));
    assert_eq!(t1.next(), None);
    assert_eq!(t2.next(), Some(3));
    assert_eq!(t2.next(), None);
    assert_eq!(t1.next(), None);
    assert_eq!(t2.next(), None);

    let (t1, t2) = xs.iter().cloned().tee();
    assert_iters_equal(t1, xs.iter().cloned());
    assert_iters_equal(t2, xs.iter().cloned());

    let (t1, t2) = xs.iter().cloned().tee();
    assert_iters_equal(t1.zip(t2), xs.iter().cloned().zip(xs.iter().cloned()));
}

#[test]
fn rciter() {
    let xs = [0, 1, 1, 1, 2, 1, 3, 5, 6];

    let mut r1 = xs.iter().cloned().into_rc();
    let mut r2 = r1.clone();
    assert_eq!(r1.next(), Some(0));
    assert_eq!(r2.next(), Some(1));
    let mut z = r1.zip(r2);
    assert_eq!(z.next(), Some((1, 1)));
    assert_eq!(z.next(), Some((2, 1)));
    assert_eq!(z.next(), Some((3, 5)));
    assert_eq!(z.next(), None);
}

#[test]
fn slice() {
    assert_iters_equal((0..10).slice(..3), 0..3);
    assert_iters_equal((0..10).slice(3..7), 3..7);
    assert_iters_equal((0..10).slice(3..27), 3..10);
    assert_iters_equal((0..10).slice(44..), 0..0);
}

#[test]
fn step() {
    assert_iters_equal((0..10).step(1), (0..10));
    assert_iters_equal((0..10).step(2), (0..10).filter(|x: &i32| *x % 2 == 0));
    assert_iters_equal((0..10).step(10), 0..1);
}

#[test]
fn trait_pointers() {
    /*
    struct ByRef<'r, I: ?Sized>(&'r mut I) where I: 'r;
    impl<'r, X, I: ?Sized> Iterator for ByRef<'r, I> where
        I: 'r + Iterator<Item=X>
    {
        type Item = X;
        fn next(&mut self) -> Option<X>
        {
            self.0.next()
        }
    }
    */
    let mut it = Box::new(0..10) as Box<Iterator<Item=i32>>;
    assert_eq!(it.next(), Some(0));

    let mut jt: &mut Iterator<Item=i32> = &mut *it;
    assert_eq!(jt.next(), Some(1));

    // they all crash/ICE at the moment.
    /*
    let mut r = ByRef(jt);
    //assert_eq!(r.next(), Some(2));

    //it.drain();
    //assert_eq!(it.dropn(2), 2);
    */
}

#[test]
fn merge() {
    assert_iters_equal((0..10).step(2).merge((1..10).step(2)), (0..10));
}

#[test]
fn to_string_join() {
    let many = [1, 2, 3];
    let one  = [1];
    let none: Vec<i32> = vec![];

    assert_eq!(many.iter().to_string_join(", "), "1, 2, 3");
    assert_eq!( one.iter().to_string_join(", "), "1");
    assert_eq!(none.iter().to_string_join(", "), "");
}

#[test]
fn multipeek() {
    let nums = vec![1u8,2,3,4,5];

    let multipeek = nums.iter().map(|&x| x).multipeek();
    assert_eq!(nums, multipeek.collect());

    let mut multipeek = nums.iter().map(|&x| x).multipeek();
    assert_eq!(multipeek.peek(), Some(&1));
    assert_eq!(multipeek.next(), Some(1));
    assert_eq!(multipeek.peek(), Some(&2));
    assert_eq!(multipeek.peek(), Some(&3));
    assert_eq!(multipeek.next(), Some(2));
    assert_eq!(multipeek.peek(), Some(&3));
    assert_eq!(multipeek.peek(), Some(&4));
    assert_eq!(multipeek.peek(), Some(&5));
    assert_eq!(multipeek.peek(), None);
    assert_eq!(multipeek.next(), Some(3));
    assert_eq!(multipeek.next(), Some(4));
    assert_eq!(multipeek.next(), Some(5));
    assert_eq!(multipeek.next(), None);
    assert_eq!(multipeek.peek(), None);

}
<commit_msg>Stop using range()<commit_after>//! Licensed under the Apache License, Version 2.0
//! http://www.apache.org/licenses/LICENSE-2.0 or the MIT license
//! http://opensource.org/licenses/MIT, at your
//! option. This file may not be copied, modified, or distributed
//! except according to those terms.

#![feature(core, test)]

#[macro_use]
extern crate itertools;

extern crate test;

use std::fmt::Debug;
use std::iter::order;
use itertools::Itertools;
use itertools::Interleave;
use itertools::Zip;

use itertools as it;

fn assert_iters_equal<A, I, J>(mut it: I, mut jt: J) where
    A: PartialEq + Debug,
    I: Iterator<Item=A>,
    J: Iterator<Item=A>,
{
    loop {
        let elti = it.next();
        let eltj = jt.next();
        assert_eq!(elti, eltj);
        if elti.is_none() { break; }
    }
}

#[test]
fn product2() {
    let s = "αβ";

    let mut prod = iproduct!(s.chars(), 0..2);
    assert!(prod.next() == Some(('α', 0)));
    assert!(prod.next() == Some(('α', 1)));
    assert!(prod.next() == Some(('β', 0)));
    assert!(prod.next() == Some(('β', 1)));
    assert!(prod.next() == None);
}

#[test]
fn product3() {
    let prod = iproduct!(0..3, 0..2, 0..2);
    assert_eq!(prod.size_hint(), (12, Some(12)));
    let v = prod.collect_vec();
    for i in 0..3 {
        for j in 0..2 {
            for k in 0..2 {
                assert!((i, j, k) == v[(i * 2 * 2 + j * 2 + k) as usize]);
            }
        }
    }
    for (_, _, _, _) in iproduct!(0..3, 0..2, 0..2, 0..3) {
        /* test compiles */
    }
}

#[test]
fn izip3() {
    let mut zip = Zip::new((0..3, 0..2, 0..2i8));
    for i in 0..2 {
        assert!((i as usize, i, i as i8) == zip.next().unwrap());
    }
    assert!(zip.next().is_none());

    let xs: [isize; 0] = [];
    let mut zip = Zip::new((0..3, 0..2, 0..2i8, xs.iter()));
    assert!(zip.next().is_none());

    for (_, _, _, _) in Zip::new((0..3, 0..2, 0..2, 0..3)) {
        /* test compiles */
    }
}

#[test]
fn write_to() {
    let xs = [7, 9, 8];
    let mut ys = [0; 5];
    let cnt = it::write(ys.iter_mut(), xs.iter().map(|x| *x));
    assert!(cnt == xs.len());
    assert!(ys == [7, 9, 8, 0, 0]);

    let cnt = it::write(ys.iter_mut(), 0..10);
    assert!(cnt == ys.len());
    assert!(ys == [0, 1, 2, 3, 4]);
}

#[test]
fn interleave() {
    let xs: [u8; 0]  = [];
    let ys = [7u8, 9, 8, 10];
    let zs = [2u8, 77];
    let it = Interleave::new(xs.iter(), ys.iter());
    assert!(order::eq(it, ys.iter()));

    let rs = [7u8, 2, 9, 77, 8, 10];
    let it = Interleave::new(ys.iter(), zs.iter());
    assert!(order::eq(it, rs.iter()));
}

#[test]
fn times() {
    assert!(it::times(0).count() == 0);
    assert!(it::times(5).count() == 5);
}

#[test]
fn drain() {
    let xs = [1i32,2,3];
    let mut sum = 0;
    xs.iter().map(|elt| sum += *elt).drain();
    assert!(sum == 6);
}

#[test]
fn foreach() {
    let xs = [1i32, 2, 3];
    let mut sum = 0;
    xs.iter().foreach(|elt| sum += *elt);
    assert!(sum == 6);
}

#[test]
fn dropn() {
    let xs = [1, 2, 3];
    let mut it = xs.iter();
    assert!(it.dropn(2) == 2);
    assert!(it.next().is_some());
    assert!(it.next().is_none());
    let mut it = xs.iter();
    assert!(it.dropn(5) == 3);
    assert!(it.next().is_none());
}

#[test]
fn dropping() {
    let xs = [1, 2, 3];
    let mut it = xs.iter().dropping(2);
    assert!(it.next().is_some());
    assert!(it.next().is_none());
    let mut it = xs.iter().dropping(5);
    assert!(it.next().is_none());
}

#[test]
fn intersperse() {
    let xs = ["a", "", "b", "c"];
    let v: Vec<&str> = xs.iter().map(|x| x.clone()).intersperse(", ").collect();
    let text: String = v.concat();
    assert_eq!(text, "a, , b, c".to_string());

    let ys = [0, 1, 2, 3];
    let mut it = ys[..0].iter().map(|x| *x).intersperse(1);
    assert!(it.next() == None);
}

#[test]
fn linspace() {
    let mut iter = it::linspace::<f32>(0., 2., 3);
    assert_eq!(iter.next(), Some(0.0));
    assert_eq!(iter.next(), Some(1.0));
    assert_eq!(iter.next(), Some(2.0));
    assert_eq!(iter.next(), None);

    let mut iter = it::linspace::<f32>(0., -2., 4);
    assert_eq!(iter.next(), Some(0.));
    assert_eq!(iter.next(), Some(-0.666666666667));
    assert_eq!(iter.next(), Some(-1.333333333333));
    assert_eq!(iter.next(), Some(-2.));
    assert_eq!(iter.next(), None);

    let mut iter = it::linspace::<f32>(0., 1., 1);
    assert_eq!(iter.next(), Some(0.));
    assert_eq!(iter.next(), None);

    let mut iter = it::linspace::<f32>(0., 1., 0);
    assert_eq!(iter.next(), None);
}

#[test]
fn dedup() {
    let xs = [0, 1, 1, 1, 2, 1, 3, 3];
    let ys = [0, 1, 2, 1, 3];
    assert_iters_equal(ys.iter(), xs.iter().dedup());
    let xs = [0, 0, 0, 0, 0];
    let ys = [0];
    assert_iters_equal(ys.iter(), xs.iter().dedup());
}

#[test]
fn batching() {
    let xs = [0, 1, 2, 1, 3];
    let ys = [(0, 1), (2, 1)];

    // An iterator that gathers elements up in pairs
    let pit = xs.iter().cloned().batching(|mut it| {
               match it.next() {
                   None => None,
                   Some(x) => match it.next() {
                       None => None,
                       Some(y) => Some((x, y)),
                   }
               }
           });
    assert_iters_equal(pit, ys.iter().cloned());
}

#[test]
fn group_by() {
    let xs = [0, 1, 1, 1, 2, 1, 3, 3];
    let ans = vec![(0, vec![0]), (1, vec![1, 1, 1]),
                   (2, vec![2]), (1, vec![1]), (3, vec![3, 3])];

    let gb = xs.iter().cloned().group_by(|elt| *elt);
    assert_iters_equal(gb, ans.into_iter());
}

#[test]
fn put_back() {
    let xs = [0, 1, 1, 1, 2, 1, 3, 3];
    let mut pb = it::PutBack::new(xs.iter().cloned());
    pb.next();
    pb.put_back(1);
    pb.put_back(0);
    assert_iters_equal(pb, xs.iter().cloned());
}

#[test]
fn tee() {
    let xs  = [0, 1, 2, 3];
    let (mut t1, mut t2) = xs.iter().cloned().tee();
    assert_eq!(t1.next(), Some(0));
    assert_eq!(t2.next(), Some(0));
    assert_eq!(t1.next(), Some(1));
    assert_eq!(t1.next(), Some(2));
    assert_eq!(t1.next(), Some(3));
    assert_eq!(t1.next(), None);
    assert_eq!(t2.next(), Some(1));
    assert_eq!(t2.next(), Some(2));
    assert_eq!(t1.next(), None);
    assert_eq!(t2.next(), Some(3));
    assert_eq!(t2.next(), None);
    assert_eq!(t1.next(), None);
    assert_eq!(t2.next(), None);

    let (t1, t2) = xs.iter().cloned().tee();
    assert_iters_equal(t1, xs.iter().cloned());
    assert_iters_equal(t2, xs.iter().cloned());

    let (t1, t2) = xs.iter().cloned().tee();
    assert_iters_equal(t1.zip(t2), xs.iter().cloned().zip(xs.iter().cloned()));
}

#[test]
fn rciter() {
    let xs = [0, 1, 1, 1, 2, 1, 3, 5, 6];

    let mut r1 = xs.iter().cloned().into_rc();
    let mut r2 = r1.clone();
    assert_eq!(r1.next(), Some(0));
    assert_eq!(r2.next(), Some(1));
    let mut z = r1.zip(r2);
    assert_eq!(z.next(), Some((1, 1)));
    assert_eq!(z.next(), Some((2, 1)));
    assert_eq!(z.next(), Some((3, 5)));
    assert_eq!(z.next(), None);
}

#[test]
fn slice() {
    assert_iters_equal((0..10).slice(..3), 0..3);
    assert_iters_equal((0..10).slice(3..7), 3..7);
    assert_iters_equal((0..10).slice(3..27), 3..10);
    assert_iters_equal((0..10).slice(44..), 0..0);
}

#[test]
fn step() {
    assert_iters_equal((0..10).step(1), (0..10));
    assert_iters_equal((0..10).step(2), (0..10).filter(|x: &i32| *x % 2 == 0));
    assert_iters_equal((0..10).step(10), 0..1);
}

#[test]
fn trait_pointers() {
    /*
    struct ByRef<'r, I: ?Sized>(&'r mut I) where I: 'r;
    impl<'r, X, I: ?Sized> Iterator for ByRef<'r, I> where
        I: 'r + Iterator<Item=X>
    {
        type Item = X;
        fn next(&mut self) -> Option<X>
        {
            self.0.next()
        }
    }
    */
    let mut it = Box::new(0..10) as Box<Iterator<Item=i32>>;
    assert_eq!(it.next(), Some(0));

    let mut jt: &mut Iterator<Item=i32> = &mut *it;
    assert_eq!(jt.next(), Some(1));

    // they all crash/ICE at the moment.
    /*
    let mut r = ByRef(jt);
    //assert_eq!(r.next(), Some(2));

    //it.drain();
    //assert_eq!(it.dropn(2), 2);
    */
}

#[test]
fn merge() {
    assert_iters_equal((0..10).step(2).merge((1..10).step(2)), (0..10));
}

#[test]
fn to_string_join() {
    let many = [1, 2, 3];
    let one  = [1];
    let none: Vec<i32> = vec![];

    assert_eq!(many.iter().to_string_join(", "), "1, 2, 3");
    assert_eq!( one.iter().to_string_join(", "), "1");
    assert_eq!(none.iter().to_string_join(", "), "");
}

#[test]
fn multipeek() {
    let nums = vec![1u8,2,3,4,5];

    let multipeek = nums.iter().map(|&x| x).multipeek();
    assert_eq!(nums, multipeek.collect());

    let mut multipeek = nums.iter().map(|&x| x).multipeek();
    assert_eq!(multipeek.peek(), Some(&1));
    assert_eq!(multipeek.next(), Some(1));
    assert_eq!(multipeek.peek(), Some(&2));
    assert_eq!(multipeek.peek(), Some(&3));
    assert_eq!(multipeek.next(), Some(2));
    assert_eq!(multipeek.peek(), Some(&3));
    assert_eq!(multipeek.peek(), Some(&4));
    assert_eq!(multipeek.peek(), Some(&5));
    assert_eq!(multipeek.peek(), None);
    assert_eq!(multipeek.next(), Some(3));
    assert_eq!(multipeek.next(), Some(4));
    assert_eq!(multipeek.next(), Some(5));
    assert_eq!(multipeek.next(), None);
    assert_eq!(multipeek.peek(), None);

}
<|endoftext|>


<commit_before>use byteorder::{LittleEndian, WriteBytesExt};
use std::io::{self, Write};

use crate::color;
use crate::error::{ImageError, ImageResult, ParameterError, ParameterErrorKind};
use crate::image::ImageEncoder;

const BITMAPFILEHEADER_SIZE: u32 = 14;
const BITMAPINFOHEADER_SIZE: u32 = 40;
const BITMAPV4HEADER_SIZE: u32 = 108;

/// The representation of a BMP encoder.
pub struct BmpEncoder<'a, W: 'a> {
    writer: &'a mut W,
}

/// BMP Encoder
///
/// An alias of [`BmpEncoder`].
///
/// TODO: remove
///
/// [`BmpEncoder`]: struct.BmpEncoder.html
#[allow(dead_code)]
#[deprecated(note = "Use `BmpEncoder` instead")]
pub type BMPEncoder<'a, W> = BmpEncoder<'a, W>;

impl<'a, W: Write + 'a> BmpEncoder<'a, W> {
    /// Create a new encoder that writes its output to ```w```.
    pub fn new(w: &'a mut W) -> Self {
        BmpEncoder { writer: w }
    }

    /// Encodes the image ```image```
    /// that has dimensions ```width``` and ```height```
    /// and ```ColorType``` ```c```.
    pub fn encode(
        &mut self,
        image: &[u8],
        width: u32,
        height: u32,
        c: color::ColorType,
    ) -> ImageResult<()> {
        self.encode_with_palette(image, width, height, c, None)
    }

    /// Same as ```encode```, but allow a palette to be passed in
    pub fn encode_with_palette(
        &mut self,
        image: &[u8],
        width: u32,
        height: u32,
        c: color::ColorType,
        palette: Option<&[(u8, u8, u8)]>,
    ) -> ImageResult<()> {
        let bmp_header_size = BITMAPFILEHEADER_SIZE;

        let (dib_header_size, written_pixel_size, palette_color_count) = get_pixel_info(c, palette)?;
        let row_pad_size = (4 - (width * written_pixel_size) % 4) % 4; // each row must be padded to a multiple of 4 bytes
        let image_size = width
            .checked_mul(height)
            .and_then(|v| v.checked_mul(written_pixel_size))
            .and_then(|v| v.checked_add(height * row_pad_size))
            .ok_or_else(|| {
                ImageError::Parameter(ParameterError::from_kind(
                    ParameterErrorKind::DimensionMismatch,
                ))
            })?;
        let palette_size = palette_color_count * 4; // all palette colors are BGRA
        let file_size = bmp_header_size + dib_header_size + palette_size + image_size;

        // write BMP header
        self.writer.write_u8(b'B')?;
        self.writer.write_u8(b'M')?;
        self.writer.write_u32::<LittleEndian>(file_size)?; // file size
        self.writer.write_u16::<LittleEndian>(0)?; // reserved 1
        self.writer.write_u16::<LittleEndian>(0)?; // reserved 2
        self.writer
            .write_u32::<LittleEndian>(bmp_header_size + dib_header_size + palette_size)?; // image data offset

        // write DIB header
        self.writer.write_u32::<LittleEndian>(dib_header_size)?;
        self.writer.write_i32::<LittleEndian>(width as i32)?;
        self.writer.write_i32::<LittleEndian>(height as i32)?;
        self.writer.write_u16::<LittleEndian>(1)?; // color planes
        self.writer
            .write_u16::<LittleEndian>((written_pixel_size * 8) as u16)?; // bits per pixel
        if dib_header_size >= BITMAPV4HEADER_SIZE {
            // Assume BGRA32
            self.writer.write_u32::<LittleEndian>(3)?; // compression method - bitfields
        } else {
            self.writer.write_u32::<LittleEndian>(0)?; // compression method - no compression
        }
        self.writer.write_u32::<LittleEndian>(image_size)?;
        self.writer.write_i32::<LittleEndian>(0)?; // horizontal ppm
        self.writer.write_i32::<LittleEndian>(0)?; // vertical ppm
        self.writer.write_u32::<LittleEndian>(palette_color_count)?;
        self.writer.write_u32::<LittleEndian>(0)?; // all colors are important
        if dib_header_size >= BITMAPV4HEADER_SIZE {
            // Assume BGRA32
            self.writer.write_u32::<LittleEndian>(0xff << 16)?; // red mask
            self.writer.write_u32::<LittleEndian>(0xff << 8)?; // green mask
            self.writer.write_u32::<LittleEndian>(0xff)?; // blue mask
            self.writer.write_u32::<LittleEndian>(0xff << 24)?; // alpha mask
            self.writer.write_u32::<LittleEndian>(0x73524742)?; // colorspace - sRGB

            // endpoints (3x3) and gamma (3)
            for _ in 0..12 {
                self.writer.write_u32::<LittleEndian>(0)?;
            }
        }

        // write image data
        match c {
            color::ColorType::Rgb8 => {
                self.encode_rgb(image, width, height, row_pad_size, 3)?
            }
            color::ColorType::Rgba8 => {
                self.encode_rgba(image, width, height, row_pad_size, 4)?
            }
            color::ColorType::L8 => {
                self.encode_gray(image, width, height, row_pad_size, 1, palette)?
            }
            color::ColorType::La8 => {
                self.encode_gray(image, width, height, row_pad_size, 2, palette)?
            }
            _ => {
                return Err(ImageError::IoError(io::Error::new(
                    io::ErrorKind::InvalidInput,
                    &get_unsupported_error_message(c)[..],
                )))
            }
        }

        Ok(())
    }

    fn encode_rgb(
        &mut self,
        image: &[u8],
        width: u32,
        height: u32,
        row_pad_size: u32,
        bytes_per_pixel: u32,
    ) -> io::Result<()> {
        let width = width as usize;
        let height = height as usize;
        let x_stride = bytes_per_pixel as usize;
        let y_stride = width * x_stride;
        for row in (0..height).rev() {
            // from the bottom up
            let row_start = row * y_stride;
            for px in image[row_start..][..y_stride].chunks_exact(x_stride) {
                let r = px[0];
                let g = px[1];
                let b = px[2];
                // written as BGR
                self.writer.write_all(&[b, g, r])?;
            }
            self.write_row_pad(row_pad_size)?;
        }

        Ok(())
    }

    fn encode_rgba(
        &mut self,
        image: &[u8],
        width: u32,
        height: u32,
        row_pad_size: u32,
        bytes_per_pixel: u32,
    ) -> io::Result<()> {
        let width = width as usize;
        let height = height as usize;
        let x_stride = bytes_per_pixel as usize;
        let y_stride = width * x_stride;
        for row in (0..height).rev() {
            // from the bottom up
            let row_start = row * y_stride;
            for px in image[row_start..][..y_stride].chunks_exact(x_stride) {
                let r = px[0];
                let g = px[1];
                let b = px[2];
                let a = px[3];
                // written as BGRA
                self.writer.write_all(&[b, g, r, a])?;
            }
            self.write_row_pad(row_pad_size)?;
        }

        Ok(())
    }

    fn encode_gray(
        &mut self,
        image: &[u8],
        width: u32,
        height: u32,
        row_pad_size: u32,
        bytes_per_pixel: u32,
        palette: Option<&[(u8, u8, u8)]>,
    ) -> io::Result<()> {
        // write grayscale palette
        if let Some(palette) = palette {
            for (r, g, b) in palette {
                // each color is written as BGRA, where A is always 0
                self.writer.write_all(&[*b, *g, *r, 0])?;
            }
        } else {
            for val in 0u8..=255 {
                // each color is written as BGRA, where A is always 0 and since only grayscale is being written, B = G = R = index
                self.writer.write_all(&[val, val, val, 0])?;
            }
        }

        // write image data
        let x_stride = bytes_per_pixel;
        let y_stride = width * x_stride;
        for row in (0..height).rev() {
            // from the bottom up
            let row_start = row * y_stride;
            for col in 0..width {
                let pixel_start = (row_start + (col * x_stride)) as usize;
                // color value is equal to the palette index
                self.writer.write_u8(image[pixel_start])?;
                // alpha is never written as it's not widely supported
            }

            self.write_row_pad(row_pad_size)?;
        }

        Ok(())
    }

    fn write_row_pad(&mut self, row_pad_size: u32) -> io::Result<()> {
        for _ in 0..row_pad_size {
            self.writer.write_u8(0)?;
        }

        Ok(())
    }
}

impl<'a, W: Write> ImageEncoder for BmpEncoder<'a, W> {
    fn write_image(
        mut self,
        buf: &[u8],
        width: u32,
        height: u32,
        color_type: color::ColorType,
    ) -> ImageResult<()> {
        self.encode(buf, width, height, color_type)
    }
}

fn get_unsupported_error_message(c: color::ColorType) -> String {
    format!(
        "Unsupported color type {:?}.  Supported types: RGB(8), RGBA(8), Gray(8), GrayA(8).",
        c
    )
}

/// Returns a tuple representing: (dib header size, written pixel size, palette color count).
fn get_pixel_info(c: color::ColorType, palette: Option<&[(u8, u8, u8)]>) -> io::Result<(u32, u32, u32)> {
    let sizes = match c {
        color::ColorType::Rgb8 => (BITMAPINFOHEADER_SIZE, 3, 0),
        color::ColorType::Rgba8 => (BITMAPV4HEADER_SIZE, 4, 0),
        color::ColorType::L8 => (BITMAPINFOHEADER_SIZE, 1, palette.map(|p| p.len()).unwrap_or(256) as u32),
        color::ColorType::La8 => (BITMAPINFOHEADER_SIZE, 1, palette.map(|p| p.len()).unwrap_or(256) as u32),
        _ => {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                &get_unsupported_error_message(c)[..],
            ))
        }
    };

    Ok(sizes)
}

#[cfg(test)]
mod tests {
    use super::super::BmpDecoder;
    use super::BmpEncoder;
    use crate::color::ColorType;
    use crate::image::ImageDecoder;
    use std::io::Cursor;

    fn round_trip_image(image: &[u8], width: u32, height: u32, c: ColorType) -> Vec<u8> {
        let mut encoded_data = Vec::new();
        {
            let mut encoder = BmpEncoder::new(&mut encoded_data);
            encoder
                .encode(&image, width, height, c)
                .expect("could not encode image");
        }

        let decoder = BmpDecoder::new(Cursor::new(&encoded_data)).expect("failed to decode");

        let mut buf = vec![0; decoder.total_bytes() as usize];
        decoder.read_image(&mut buf).expect("failed to decode");
        buf
    }

    #[test]
    fn round_trip_single_pixel_rgb() {
        let image = [255u8, 0, 0]; // single red pixel
        let decoded = round_trip_image(&image, 1, 1, ColorType::Rgb8);
        assert_eq!(3, decoded.len());
        assert_eq!(255, decoded[0]);
        assert_eq!(0, decoded[1]);
        assert_eq!(0, decoded[2]);
    }

    #[test]
    fn huge_files_return_error() {
        let mut encoded_data = Vec::new();
        let image = vec![0u8; 3 * 40_000 * 40_000]; // 40_000x40_000 pixels, 3 bytes per pixel, allocated on the heap
        let mut encoder = BmpEncoder::new(&mut encoded_data);
        let result = encoder.encode(&image, 40_000, 40_000, ColorType::Rgb8);
        assert!(result.is_err());
    }

    #[test]
    fn round_trip_single_pixel_rgba() {
        let image = [1, 2, 3, 4];
        let decoded = round_trip_image(&image, 1, 1, ColorType::Rgba8);
        assert_eq!(&decoded[..], &image[..]);
    }

    #[test]
    fn round_trip_3px_rgb() {
        let image = [0u8; 3 * 3 * 3]; // 3x3 pixels, 3 bytes per pixel
        let _decoded = round_trip_image(&image, 3, 3, ColorType::Rgb8);
    }

    #[test]
    fn round_trip_gray() {
        let image = [0u8, 1, 2]; // 3 pixels
        let decoded = round_trip_image(&image, 3, 1, ColorType::L8);
        // should be read back as 3 RGB pixels
        assert_eq!(9, decoded.len());
        assert_eq!(0, decoded[0]);
        assert_eq!(0, decoded[1]);
        assert_eq!(0, decoded[2]);
        assert_eq!(1, decoded[3]);
        assert_eq!(1, decoded[4]);
        assert_eq!(1, decoded[5]);
        assert_eq!(2, decoded[6]);
        assert_eq!(2, decoded[7]);
        assert_eq!(2, decoded[8]);
    }

    #[test]
    fn round_trip_graya() {
        let image = [0u8, 0, 1, 0, 2, 0]; // 3 pixels, each with an alpha channel
        let decoded = round_trip_image(&image, 1, 3, ColorType::La8);
        // should be read back as 3 RGB pixels
        assert_eq!(9, decoded.len());
        assert_eq!(0, decoded[0]);
        assert_eq!(0, decoded[1]);
        assert_eq!(0, decoded[2]);
        assert_eq!(1, decoded[3]);
        assert_eq!(1, decoded[4]);
        assert_eq!(1, decoded[5]);
        assert_eq!(2, decoded[6]);
        assert_eq!(2, decoded[7]);
        assert_eq!(2, decoded[8]);
    }
}
<commit_msg>(u8,u8,u8) -> [u8;3] on encoder<commit_after>use byteorder::{LittleEndian, WriteBytesExt};
use std::io::{self, Write};

use crate::color;
use crate::error::{ImageError, ImageResult, ParameterError, ParameterErrorKind};
use crate::image::ImageEncoder;

const BITMAPFILEHEADER_SIZE: u32 = 14;
const BITMAPINFOHEADER_SIZE: u32 = 40;
const BITMAPV4HEADER_SIZE: u32 = 108;

/// The representation of a BMP encoder.
pub struct BmpEncoder<'a, W: 'a> {
    writer: &'a mut W,
}

/// BMP Encoder
///
/// An alias of [`BmpEncoder`].
///
/// TODO: remove
///
/// [`BmpEncoder`]: struct.BmpEncoder.html
#[allow(dead_code)]
#[deprecated(note = "Use `BmpEncoder` instead")]
pub type BMPEncoder<'a, W> = BmpEncoder<'a, W>;

impl<'a, W: Write + 'a> BmpEncoder<'a, W> {
    /// Create a new encoder that writes its output to ```w```.
    pub fn new(w: &'a mut W) -> Self {
        BmpEncoder { writer: w }
    }

    /// Encodes the image ```image```
    /// that has dimensions ```width``` and ```height```
    /// and ```ColorType``` ```c```.
    pub fn encode(
        &mut self,
        image: &[u8],
        width: u32,
        height: u32,
        c: color::ColorType,
    ) -> ImageResult<()> {
        self.encode_with_palette(image, width, height, c, None)
    }

    /// Same as ```encode```, but allow a palette to be passed in
    pub fn encode_with_palette(
        &mut self,
        image: &[u8],
        width: u32,
        height: u32,
        c: color::ColorType,
        palette: Option<&[[u8; 3]]>,
    ) -> ImageResult<()> {
        let bmp_header_size = BITMAPFILEHEADER_SIZE;

        let (dib_header_size, written_pixel_size, palette_color_count) = get_pixel_info(c, palette)?;
        let row_pad_size = (4 - (width * written_pixel_size) % 4) % 4; // each row must be padded to a multiple of 4 bytes
        let image_size = width
            .checked_mul(height)
            .and_then(|v| v.checked_mul(written_pixel_size))
            .and_then(|v| v.checked_add(height * row_pad_size))
            .ok_or_else(|| {
                ImageError::Parameter(ParameterError::from_kind(
                    ParameterErrorKind::DimensionMismatch,
                ))
            })?;
        let palette_size = palette_color_count * 4; // all palette colors are BGRA
        let file_size = bmp_header_size + dib_header_size + palette_size + image_size;

        // write BMP header
        self.writer.write_u8(b'B')?;
        self.writer.write_u8(b'M')?;
        self.writer.write_u32::<LittleEndian>(file_size)?; // file size
        self.writer.write_u16::<LittleEndian>(0)?; // reserved 1
        self.writer.write_u16::<LittleEndian>(0)?; // reserved 2
        self.writer
            .write_u32::<LittleEndian>(bmp_header_size + dib_header_size + palette_size)?; // image data offset

        // write DIB header
        self.writer.write_u32::<LittleEndian>(dib_header_size)?;
        self.writer.write_i32::<LittleEndian>(width as i32)?;
        self.writer.write_i32::<LittleEndian>(height as i32)?;
        self.writer.write_u16::<LittleEndian>(1)?; // color planes
        self.writer
            .write_u16::<LittleEndian>((written_pixel_size * 8) as u16)?; // bits per pixel
        if dib_header_size >= BITMAPV4HEADER_SIZE {
            // Assume BGRA32
            self.writer.write_u32::<LittleEndian>(3)?; // compression method - bitfields
        } else {
            self.writer.write_u32::<LittleEndian>(0)?; // compression method - no compression
        }
        self.writer.write_u32::<LittleEndian>(image_size)?;
        self.writer.write_i32::<LittleEndian>(0)?; // horizontal ppm
        self.writer.write_i32::<LittleEndian>(0)?; // vertical ppm
        self.writer.write_u32::<LittleEndian>(palette_color_count)?;
        self.writer.write_u32::<LittleEndian>(0)?; // all colors are important
        if dib_header_size >= BITMAPV4HEADER_SIZE {
            // Assume BGRA32
            self.writer.write_u32::<LittleEndian>(0xff << 16)?; // red mask
            self.writer.write_u32::<LittleEndian>(0xff << 8)?; // green mask
            self.writer.write_u32::<LittleEndian>(0xff)?; // blue mask
            self.writer.write_u32::<LittleEndian>(0xff << 24)?; // alpha mask
            self.writer.write_u32::<LittleEndian>(0x73524742)?; // colorspace - sRGB

            // endpoints (3x3) and gamma (3)
            for _ in 0..12 {
                self.writer.write_u32::<LittleEndian>(0)?;
            }
        }

        // write image data
        match c {
            color::ColorType::Rgb8 => {
                self.encode_rgb(image, width, height, row_pad_size, 3)?
            }
            color::ColorType::Rgba8 => {
                self.encode_rgba(image, width, height, row_pad_size, 4)?
            }
            color::ColorType::L8 => {
                self.encode_gray(image, width, height, row_pad_size, 1, palette)?
            }
            color::ColorType::La8 => {
                self.encode_gray(image, width, height, row_pad_size, 2, palette)?
            }
            _ => {
                return Err(ImageError::IoError(io::Error::new(
                    io::ErrorKind::InvalidInput,
                    &get_unsupported_error_message(c)[..],
                )))
            }
        }

        Ok(())
    }

    fn encode_rgb(
        &mut self,
        image: &[u8],
        width: u32,
        height: u32,
        row_pad_size: u32,
        bytes_per_pixel: u32,
    ) -> io::Result<()> {
        let width = width as usize;
        let height = height as usize;
        let x_stride = bytes_per_pixel as usize;
        let y_stride = width * x_stride;
        for row in (0..height).rev() {
            // from the bottom up
            let row_start = row * y_stride;
            for px in image[row_start..][..y_stride].chunks_exact(x_stride) {
                let r = px[0];
                let g = px[1];
                let b = px[2];
                // written as BGR
                self.writer.write_all(&[b, g, r])?;
            }
            self.write_row_pad(row_pad_size)?;
        }

        Ok(())
    }

    fn encode_rgba(
        &mut self,
        image: &[u8],
        width: u32,
        height: u32,
        row_pad_size: u32,
        bytes_per_pixel: u32,
    ) -> io::Result<()> {
        let width = width as usize;
        let height = height as usize;
        let x_stride = bytes_per_pixel as usize;
        let y_stride = width * x_stride;
        for row in (0..height).rev() {
            // from the bottom up
            let row_start = row * y_stride;
            for px in image[row_start..][..y_stride].chunks_exact(x_stride) {
                let r = px[0];
                let g = px[1];
                let b = px[2];
                let a = px[3];
                // written as BGRA
                self.writer.write_all(&[b, g, r, a])?;
            }
            self.write_row_pad(row_pad_size)?;
        }

        Ok(())
    }

    fn encode_gray(
        &mut self,
        image: &[u8],
        width: u32,
        height: u32,
        row_pad_size: u32,
        bytes_per_pixel: u32,
        palette: Option<&[[u8; 3]]>,
    ) -> io::Result<()> {
        // write grayscale palette
        if let Some(palette) = palette {
            for item in palette {
                // each color is written as BGRA, where A is always 0
                self.writer.write_all(&[item[2], item[1], item[0], 0])?;
            }
        } else {
            for val in 0u8..=255 {
                // each color is written as BGRA, where A is always 0 and since only grayscale is being written, B = G = R = index
                self.writer.write_all(&[val, val, val, 0])?;
            }
        }

        // write image data
        let x_stride = bytes_per_pixel;
        let y_stride = width * x_stride;
        for row in (0..height).rev() {
            // from the bottom up
            let row_start = row * y_stride;
            for col in 0..width {
                let pixel_start = (row_start + (col * x_stride)) as usize;
                // color value is equal to the palette index
                self.writer.write_u8(image[pixel_start])?;
                // alpha is never written as it's not widely supported
            }

            self.write_row_pad(row_pad_size)?;
        }

        Ok(())
    }

    fn write_row_pad(&mut self, row_pad_size: u32) -> io::Result<()> {
        for _ in 0..row_pad_size {
            self.writer.write_u8(0)?;
        }

        Ok(())
    }
}

impl<'a, W: Write> ImageEncoder for BmpEncoder<'a, W> {
    fn write_image(
        mut self,
        buf: &[u8],
        width: u32,
        height: u32,
        color_type: color::ColorType,
    ) -> ImageResult<()> {
        self.encode(buf, width, height, color_type)
    }
}

fn get_unsupported_error_message(c: color::ColorType) -> String {
    format!(
        "Unsupported color type {:?}.  Supported types: RGB(8), RGBA(8), Gray(8), GrayA(8).",
        c
    )
}

/// Returns a tuple representing: (dib header size, written pixel size, palette color count).
fn get_pixel_info(c: color::ColorType, palette: Option<&[[u8; 3]]>) -> io::Result<(u32, u32, u32)> {
    let sizes = match c {
        color::ColorType::Rgb8 => (BITMAPINFOHEADER_SIZE, 3, 0),
        color::ColorType::Rgba8 => (BITMAPV4HEADER_SIZE, 4, 0),
        color::ColorType::L8 => (BITMAPINFOHEADER_SIZE, 1, palette.map(|p| p.len()).unwrap_or(256) as u32),
        color::ColorType::La8 => (BITMAPINFOHEADER_SIZE, 1, palette.map(|p| p.len()).unwrap_or(256) as u32),
        _ => {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                &get_unsupported_error_message(c)[..],
            ))
        }
    };

    Ok(sizes)
}

#[cfg(test)]
mod tests {
    use super::super::BmpDecoder;
    use super::BmpEncoder;
    use crate::color::ColorType;
    use crate::image::ImageDecoder;
    use std::io::Cursor;

    fn round_trip_image(image: &[u8], width: u32, height: u32, c: ColorType) -> Vec<u8> {
        let mut encoded_data = Vec::new();
        {
            let mut encoder = BmpEncoder::new(&mut encoded_data);
            encoder
                .encode(&image, width, height, c)
                .expect("could not encode image");
        }

        let decoder = BmpDecoder::new(Cursor::new(&encoded_data)).expect("failed to decode");

        let mut buf = vec![0; decoder.total_bytes() as usize];
        decoder.read_image(&mut buf).expect("failed to decode");
        buf
    }

    #[test]
    fn round_trip_single_pixel_rgb() {
        let image = [255u8, 0, 0]; // single red pixel
        let decoded = round_trip_image(&image, 1, 1, ColorType::Rgb8);
        assert_eq!(3, decoded.len());
        assert_eq!(255, decoded[0]);
        assert_eq!(0, decoded[1]);
        assert_eq!(0, decoded[2]);
    }

    #[test]
    fn huge_files_return_error() {
        let mut encoded_data = Vec::new();
        let image = vec![0u8; 3 * 40_000 * 40_000]; // 40_000x40_000 pixels, 3 bytes per pixel, allocated on the heap
        let mut encoder = BmpEncoder::new(&mut encoded_data);
        let result = encoder.encode(&image, 40_000, 40_000, ColorType::Rgb8);
        assert!(result.is_err());
    }

    #[test]
    fn round_trip_single_pixel_rgba() {
        let image = [1, 2, 3, 4];
        let decoded = round_trip_image(&image, 1, 1, ColorType::Rgba8);
        assert_eq!(&decoded[..], &image[..]);
    }

    #[test]
    fn round_trip_3px_rgb() {
        let image = [0u8; 3 * 3 * 3]; // 3x3 pixels, 3 bytes per pixel
        let _decoded = round_trip_image(&image, 3, 3, ColorType::Rgb8);
    }

    #[test]
    fn round_trip_gray() {
        let image = [0u8, 1, 2]; // 3 pixels
        let decoded = round_trip_image(&image, 3, 1, ColorType::L8);
        // should be read back as 3 RGB pixels
        assert_eq!(9, decoded.len());
        assert_eq!(0, decoded[0]);
        assert_eq!(0, decoded[1]);
        assert_eq!(0, decoded[2]);
        assert_eq!(1, decoded[3]);
        assert_eq!(1, decoded[4]);
        assert_eq!(1, decoded[5]);
        assert_eq!(2, decoded[6]);
        assert_eq!(2, decoded[7]);
        assert_eq!(2, decoded[8]);
    }

    #[test]
    fn round_trip_graya() {
        let image = [0u8, 0, 1, 0, 2, 0]; // 3 pixels, each with an alpha channel
        let decoded = round_trip_image(&image, 1, 3, ColorType::La8);
        // should be read back as 3 RGB pixels
        assert_eq!(9, decoded.len());
        assert_eq!(0, decoded[0]);
        assert_eq!(0, decoded[1]);
        assert_eq!(0, decoded[2]);
        assert_eq!(1, decoded[3]);
        assert_eq!(1, decoded[4]);
        assert_eq!(1, decoded[5]);
        assert_eq!(2, decoded[6]);
        assert_eq!(2, decoded[7]);
        assert_eq!(2, decoded[8]);
    }
}
<|endoftext|>


<commit_before>
/// Internal macro to provide speedup for `FromIterator::from_iter`.
///
/// Needed as of rustc 1.7.0-nightly (81ae8be71 2015-12-09).
macro_rules! run_from_iter(
    ( $iter:ident as $r_ty:ty ) => {
        {
            // Dummy variable to tie the scope of the iteration into the current scope,
            // this causes rustc to inline `Iterator::next` into this scope including
            // `FromIterator::from_iter`.
            // This is probably how it works, #[inline] and #[inline(always)] do not affect it as
            // of rustc 1.7.0-nightly (81ae8be71 2015-12-09).
            let mut item = false;

            // the inspect() is useless here, but ties the inner scope of the loop to the current
            // scope which will make it inline `Iterator::next`.
            let result: T = FromIterator::from_iter($iter.by_ref().inspect(|_| item = true));

            // Oddly enough, no inspect() is much much slower:
            //
            // ```
            // let result: $r_ty = FromIterator::from_iter($iter.by_ref());
            // ```
            //
            // The above code is only faster in very very small benchmarks, anything larger than
            // the benchmarks for many (ie. multiple `many(i, any)`) will most likely be slower
            // when using the above line.

            // The performance difference is definitely noticeable in benchmarks, parsing 10k bytes of
            // any type and just storing them in a Vec is 10x slower on current nightly
            // (rustc 1.7.0-nightly (81ae8be71 2015-12-09)).

            (result, $iter.end_state())
        }
    }
);

/// Macro to implement and run a parser iterator, it provides the ability to add an extra state
/// variable into it and also provide a size_hint as well as a pre- and on-next hooks.
macro_rules! run_iter {
    (
        input:     $input:expr,
        parser:    $parser:expr,
        state:     $data_ty:ty : $data:expr,

        size_hint($size_hint_self:ident) $size_hint:block
        next($next_self:ident) {
            pre $pre_next:block
            on  $on_next:block
        }

        => $result:ident : $t:ty {
             $($pat:pat => $arm:expr),*
        }
    ) => { {
        enum EndState<'a, I, E>
          where I: 'a {
            Error(&'a [I], E),
            Incomplete(usize),
        }

        struct Iter<'a, I, T, E, F>
          where I: 'a,
                T: 'a,
                E: 'a,
                F: FnMut(Input<'a, I>) -> ParseResult<'a, I, T, E> {
            /// Last state of the parser
            state:  EndState<'a, I, E>,
            /// Parser to execute once for each iteration
            parser: F,
            /// Remaining buffer
            buf:    Input<'a, I>,
            /// Nested state
            data:   $data_ty,
            _t:     PhantomData<T>,
        }

        impl<'a, I, T, E, F> Iter<'a, I, T, E, F>
          where I: 'a,
                T: 'a,
                E: 'a,
                F: FnMut(Input<'a, I>) -> ParseResult<'a, I, T, E> {
            #[inline]
            fn end_state(self) -> (Input<'a, I>, $data_ty, EndState<'a, I, E>) {
                (self.buf, self.data, self.state)
            }
        }

        impl<'a, I, T, E, F> Iterator for Iter<'a, I, T, E, F>
          where I: 'a,
                T: 'a,
                E: 'a,
                F: FnMut(Input<'a, I>) -> ParseResult<'a, I, T, E> {
            type Item = T;

            #[inline]
            fn size_hint(&$size_hint_self) -> (usize, Option<usize>) {
                $size_hint
            }

            #[inline]
            fn next(&mut $next_self) -> Option<Self::Item> {
                $pre_next

                match ($next_self.parser)($next_self.buf.clone()).into_inner() {
                    State::Data(b, v) => {
                        $next_self.buf    = b;

                        $on_next

                        Some(v)
                    },
                    State::Error(b, e) => {
                        $next_self.state = EndState::Error(b, e);

                        None
                    },
                    State::Incomplete(n) => {
                        $next_self.state = EndState::Incomplete(n);

                        None
                    },
                }
            }
        }

        let mut iter = Iter {
            state:  EndState::Incomplete(1),
            parser: $parser,
            buf:    $input,
            data:   $data,
            _t:     PhantomData,
        };

        let ($result, state) = run_from_iter!(iter as $t);

        match state {
            $($pat => $arm),*
        }
    } }
}

/// Version of run_iter which allows for an additional parser to be run which can terminate
/// iteration early.
macro_rules! run_iter_till {
    (
        input:     $input:expr,
        parser:    $parser:expr,
        end:       $end:expr,
        state:     $data_ty:ty : $data:expr,

        size_hint($size_hint_self:ident) $size_hint:block
        next($next_self:ident) {
            pre $pre_next:block
            on  $on_next:block
        }

        => $result:ident : $t:ty {
             $($pat:pat => $arm:expr),*
        }
    ) => { {
        enum EndStateTill<'a, I, E>
          where I: 'a {
            Error(&'a [I], E),
            Incomplete(usize),
            EndSuccess,
        }

        /// Iterator used by ``many_till`` and ``many1``.
        struct IterTill<'a, I, T, U, E, F, P, N>
          where I: 'a,
                T: 'a,
                E: 'a,
                U: 'a,
                N: 'a,
                P: FnMut(Input<'a, I>) -> ParseResult<'a, I, T, E>,
                F: FnMut(Input<'a, I>) -> ParseResult<'a, I, U, N> {
            state:  EndStateTill<'a, I, E>,
            parser: P,
            end:    F,
            buf:    Input<'a, I>,
            data:   $data_ty,
            _t:     PhantomData<(T, U, N)>,
        }

        impl<'a, I, T, U, E, F, P, N> IterTill<'a, I, T, U, E, F, P, N>
          where I: 'a,
                T: 'a,
                E: 'a,
                U: 'a,
                N: 'a,
                P: FnMut(Input<'a, I>) -> ParseResult<'a, I, T, E>,
                F: FnMut(Input<'a, I>) -> ParseResult<'a, I, U, N> {
            /// Destructures the iterator returning the position just after the last successful parse as
            /// well as the state of the last attempt to parse data.
            #[inline]
            fn end_state(self) -> (Input<'a, I>, $data_ty, EndStateTill<'a, I, E>) {
                (self.buf, self.data, self.state)
            }
        }

        impl<'a, I, T, U, E, F, P, N> Iterator for IterTill<'a, I, T, U, E, F, P, N>
          where I: 'a,
                T: 'a,
                E: 'a,
                U: 'a,
                N: 'a,
                P: FnMut(Input<'a, I>) -> ParseResult<'a, I, T, E>,
                F: FnMut(Input<'a, I>) -> ParseResult<'a, I, U, N> {
            type Item = T;

            #[inline]
            fn size_hint(&$size_hint_self) -> (usize, Option<usize>) {
                $size_hint
            }

            #[inline]
            fn next(&mut $next_self) -> Option<Self::Item> {
                $pre_next

                match ($next_self.parser)($next_self.buf.clone()).into_inner() {
                    State::Data(b, v) => {
                        $next_self.buf = b;

                        $on_next

                        Some(v)
                    },
                    State::Error(b, e) => {
                        $next_self.state = EndStateTill::Error(b, e);

                        None
                    },
                    State::Incomplete(n) => {
                        $next_self.state = EndStateTill::Incomplete(n);

                        None
                    },
                }
            }
        }

        let mut iter = IterTill {
            state:  EndStateTill::Incomplete(1),
            parser: $parser,
            end:    $end,
            buf:    $input,
            data:   $data,
            _t:     PhantomData,
        };

        let ($result, state) = run_from_iter!(iter as $t);

        match state {
            $($pat => $arm),*
        }
    } }
}

/// Used with `run_iter_till!` macro to attempt to end iteration early. If the test succeeds the
/// buffer position will be updated and the state set to `EndStateTill::EndSuccess` and a `None`
/// will be returned, stopping the iteration. If the test fails execution continues.
macro_rules! iter_till_end_test {
    ( $the_self:ident ) => { {
        if let State::Data(b, _) = ($the_self.end)($the_self.buf.clone()).into_inner() {
            $the_self.buf   = b;
            $the_self.state = EndStateTill::EndSuccess;

            return None
        }
    } }
}
<commit_msg>perf/combinators: Removed inspect-hack which encouraged inlining of iterator<commit_after>
/// Macro to implement and run a parser iterator, it provides the ability to add an extra state
/// variable into it and also provide a size_hint as well as a pre- and on-next hooks.
macro_rules! run_iter {
    (
        input:     $input:expr,
        parser:    $parser:expr,
        state:     $data_ty:ty : $data:expr,

        size_hint($size_hint_self:ident) $size_hint:block
        next($next_self:ident) {
            pre $pre_next:block
            on  $on_next:block
        }

        => $result:ident : $t:ty {
             $($pat:pat => $arm:expr),*
        }
    ) => { {
        enum EndState<'a, I, E>
          where I: 'a {
            Error(&'a [I], E),
            Incomplete(usize),
        }

        struct Iter<'a, I, T, E, F>
          where I: 'a,
                T: 'a,
                E: 'a,
                F: FnMut(Input<'a, I>) -> ParseResult<'a, I, T, E> {
            /// Last state of the parser
            state:  EndState<'a, I, E>,
            /// Parser to execute once for each iteration
            parser: F,
            /// Remaining buffer
            buf:    Input<'a, I>,
            /// Nested state
            data:   $data_ty,
            _t:     PhantomData<T>,
        }

        impl<'a, I, T, E, F> Iter<'a, I, T, E, F>
          where I: 'a,
                T: 'a,
                E: 'a,
                F: FnMut(Input<'a, I>) -> ParseResult<'a, I, T, E> {
            #[inline]
            fn end_state(self) -> (Input<'a, I>, $data_ty, EndState<'a, I, E>) {
                (self.buf, self.data, self.state)
            }
        }

        impl<'a, I, T, E, F> Iterator for Iter<'a, I, T, E, F>
          where I: 'a,
                T: 'a,
                E: 'a,
                F: FnMut(Input<'a, I>) -> ParseResult<'a, I, T, E> {
            type Item = T;

            #[inline]
            fn size_hint(&$size_hint_self) -> (usize, Option<usize>) {
                $size_hint
            }

            #[inline]
            fn next(&mut $next_self) -> Option<Self::Item> {
                $pre_next

                match ($next_self.parser)($next_self.buf.clone()).into_inner() {
                    State::Data(b, v) => {
                        $next_self.buf    = b;

                        $on_next

                        Some(v)
                    },
                    State::Error(b, e) => {
                        $next_self.state = EndState::Error(b, e);

                        None
                    },
                    State::Incomplete(n) => {
                        $next_self.state = EndState::Incomplete(n);

                        None
                    },
                }
            }
        }

        let mut iter = Iter {
            state:  EndState::Incomplete(1),
            parser: $parser,
            buf:    $input,
            data:   $data,
            _t:     PhantomData,
        };

        let $result: $t = FromIterator::from_iter(iter.by_ref());

        match iter.end_state() {
            $($pat => $arm),*
        }
    } }
}

/// Version of run_iter which allows for an additional parser to be run which can terminate
/// iteration early.
macro_rules! run_iter_till {
    (
        input:     $input:expr,
        parser:    $parser:expr,
        end:       $end:expr,
        state:     $data_ty:ty : $data:expr,

        size_hint($size_hint_self:ident) $size_hint:block
        next($next_self:ident) {
            pre $pre_next:block
            on  $on_next:block
        }

        => $result:ident : $t:ty {
             $($pat:pat => $arm:expr),*
        }
    ) => { {
        enum EndStateTill<'a, I, E>
          where I: 'a {
            Error(&'a [I], E),
            Incomplete(usize),
            EndSuccess,
        }

        /// Iterator used by ``many_till`` and ``many1``.
        struct IterTill<'a, I, T, U, E, F, P, N>
          where I: 'a,
                T: 'a,
                E: 'a,
                U: 'a,
                N: 'a,
                P: FnMut(Input<'a, I>) -> ParseResult<'a, I, T, E>,
                F: FnMut(Input<'a, I>) -> ParseResult<'a, I, U, N> {
            state:  EndStateTill<'a, I, E>,
            parser: P,
            end:    F,
            buf:    Input<'a, I>,
            data:   $data_ty,
            _t:     PhantomData<(T, U, N)>,
        }

        impl<'a, I, T, U, E, F, P, N> IterTill<'a, I, T, U, E, F, P, N>
          where I: 'a,
                T: 'a,
                E: 'a,
                U: 'a,
                N: 'a,
                P: FnMut(Input<'a, I>) -> ParseResult<'a, I, T, E>,
                F: FnMut(Input<'a, I>) -> ParseResult<'a, I, U, N> {
            /// Destructures the iterator returning the position just after the last successful parse as
            /// well as the state of the last attempt to parse data.
            #[inline]
            fn end_state(self) -> (Input<'a, I>, $data_ty, EndStateTill<'a, I, E>) {
                (self.buf, self.data, self.state)
            }
        }

        impl<'a, I, T, U, E, F, P, N> Iterator for IterTill<'a, I, T, U, E, F, P, N>
          where I: 'a,
                T: 'a,
                E: 'a,
                U: 'a,
                N: 'a,
                P: FnMut(Input<'a, I>) -> ParseResult<'a, I, T, E>,
                F: FnMut(Input<'a, I>) -> ParseResult<'a, I, U, N> {
            type Item = T;

            #[inline]
            fn size_hint(&$size_hint_self) -> (usize, Option<usize>) {
                $size_hint
            }

            #[inline]
            fn next(&mut $next_self) -> Option<Self::Item> {
                $pre_next

                match ($next_self.parser)($next_self.buf.clone()).into_inner() {
                    State::Data(b, v) => {
                        $next_self.buf = b;

                        $on_next

                        Some(v)
                    },
                    State::Error(b, e) => {
                        $next_self.state = EndStateTill::Error(b, e);

                        None
                    },
                    State::Incomplete(n) => {
                        $next_self.state = EndStateTill::Incomplete(n);

                        None
                    },
                }
            }
        }

        let mut iter = IterTill {
            state:  EndStateTill::Incomplete(1),
            parser: $parser,
            end:    $end,
            buf:    $input,
            data:   $data,
            _t:     PhantomData,
        };

        let $result: $t = FromIterator::from_iter(iter.by_ref());

        match iter.end_state() {
            $($pat => $arm),*
        }
    } }
}

/// Used with `run_iter_till!` macro to attempt to end iteration early. If the test succeeds the
/// buffer position will be updated and the state set to `EndStateTill::EndSuccess` and a `None`
/// will be returned, stopping the iteration. If the test fails execution continues.
macro_rules! iter_till_end_test {
    ( $the_self:ident ) => { {
        if let State::Data(b, _) = ($the_self.end)($the_self.buf.clone()).into_inner() {
            $the_self.buf   = b;
            $the_self.state = EndStateTill::EndSuccess;

            return None
        }
    } }
}
<|endoftext|>