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

#[macro_use]
extern crate log;
extern crate osmpbfreader;
extern crate rustc_serialize;
extern crate docopt;
extern crate mimir;
extern crate rs_es;
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 admin_by_level(admins: &AdminsVec, city_level: u32) -> Option<&Rc<mimir::Admin>> {
    match admins.iter().position(|adm| adm.level == city_level) {
        Some(idx) => Some(&admins[idx]),
        None => None
    }
}

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);
                    let label = match admin_by_level(&admin, city_level) {
                        Some(adm) => format!("{} ({})", way_name, adm.label),
                        None =>  way_name.to_string()
                    };
                    ret ret(street_list.push(mimir::Street {
                                id: way.id.to_string(),
                                street_name: way_name.to_string(),
                                label: label,
                                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);
            let label = match admin_by_level(&admins, city_level) {
                Some(adm) => format!("{} ({})", way_name, adm.label),
                None =>  way_name.to_string()
            };
            ret ret(street_list.push(mimir::Street{
   	                    id: way.id.to_string(),
   	                    street_name: way_name.to_string(),
   	                    label: label,
   	                    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>//! Cartridge emulation. There are multiple cartridge types with
//! different capabilities (bankable ROM/RAM, battery, RTC etc...).

use std::fmt::{Show, Formatter, Error};
use std::io::{File, Reader, Writer, IoResult, Open, ReadWrite, SeekSet};

mod models;

/// Common state for all cartridge types
pub struct Cartridge {
    /// Cartridge ROM data
    rom:        Vec<u8>,
    /// Cartridge RAM data
    ram:        Vec<u8>,
    /// Current number of the rom bank mapped at [0x4000, 0x7fff]
    rom_bank:   u8,
    /// Current bank offset for the bank mapped at [0x4000, 0x7fff].
    /// This value is added to ROM register addresses when they're in
    /// that range.
    rom_offset: uint,
    /// Current bank offset for the RAM
    ram_offset: uint,
    /// If `true` RAM is write protected
    ram_wp:     bool,
    /// Certain cartridges allow banking either the RAM or ROM
    /// depending on the value of this flag.
    bank_ram:   bool,
    /// struct used to handle model specific functions
    model:      models::Model,
    /// Path to the ROM image for this cartridge
    path:       Path,
    /// optional save file used to store non-volatile RAM on emulator
    /// shutdown
    save_file:  Option<File>,
}

impl Cartridge {
    /// Load a Cartridge ROM from `path`.
    pub fn from_path(rom_path: &Path) -> IoResult<Cartridge> {
        let mut source = try!(File::open(rom_path));

        // There must always be at least two ROM banks
        let rom = try!(source.read_exact(2 * ROM_BANK_SIZE));

        let model = models::from_id(rom[offsets::TYPE]);

        let mut cartridge = Cartridge {
            rom:        rom,
            ram:        Vec::new(),
            // Default to bank 1 for bankable region
            rom_bank:   1,
            rom_offset: 0,
            ram_offset: 0,
            ram_wp:     false,
            bank_ram:   false,
            model:      model,
            path:       rom_path.clone(),
            save_file:  None,
        };

        let rombanks = match cartridge.rom_banks() {
            Some(n) => n,
            None    => panic!("Can't determine ROM size"),
        };

        // Read the remaining roms banks
        if rombanks > 2 {
            let remb      = rombanks - 2;
            let mut off   = 2    * ROM_BANK_SIZE;
            let mut remsz = remb * ROM_BANK_SIZE;

            // Reserve space for the remaining banks
            cartridge.rom.grow(remsz, 0);

            while remsz > 0 {
                let r = try!(source.read(cartridge.rom.slice_from_mut(off)));

                remsz -= r;
                off   += r;
            }
        }

        try!(cartridge.init_ram());

        Ok(cartridge)
    }

    /// Init cartridge RAM and tie it with a `File` for saving if
    /// necessary.
    fn init_ram(&mut self) -> IoResult<()> {
        let (rambanks, banksize) = match self.ram_banks() {
            Some(v) => v,
            None    => panic!("Can't determine RAM size"),
        };

        let ramsize = rambanks * banksize;

        if ramsize == 0 {
            // No RAM on this cartridge, we're done
            return Ok(());
        }

        // We have some RAM, open the save file or create it if it
        // doesn't exist yet
        let mut savepath = self.path.clone();
        savepath.set_extension("sav");

        let mut save_file = try!(File::open_mode(&savepath,
                                                Open,
                                                ReadWrite));

        let save_size = try!(save_file.stat()).size;

        if save_size == 0 {
            // The file is empty (probably new). initialize
            // the RAM with 0s.
            self.ram.grow(ramsize, 0);
            // Then fill the file with the right amount of 0s
            // to reserve enough space for saving later.
            try!(save_file.write(self.ram.as_slice()));
        } else if save_size == (ramsize as u64) {
            // The file contains a RAM image
            self.ram = try!(save_file.read_exact(ramsize as uint));
        } else {
            panic!("Unexpected save file size for {}: expected {} got {}",
                   savepath.display(), ramsize, save_size);
        }

        // Store the file handle to save progress later
        self.save_file = Some(save_file);

        Ok(())
    }

    /// Update the save file
    pub fn save_ram(&mut self) -> IoResult<()> {
        if let Some(mut f) = self.save_file.as_mut() {
            // Rewind to the beginning of the file and update its
            // contents
            println!("Saving to {}", f.path().display());

            try!(f.seek(0, SeekSet));
            try!(f.write(self.ram.as_slice()));
        }

        Ok(())
    }

    /// Attempt to retreive the rom's name
    pub fn name(&self) -> Option<String> {
        let mut name = String::with_capacity(16);

        for i in range(0, 16) {
            let c = self.rom[offsets::TITLE + i].to_ascii();

            // If the name is shorter than 16bytes it's padded with 0s
            if c == 0.to_ascii() {
                break;
            }

            // Only uppercase ASCII is valid
            if !(c.is_uppercase() || c.is_blank() || c.is_punctuation()) {
                return None;
            }

            // Append new character
            name.grow(1, c.as_char());
        }

        Some(name)
    }

    /// Return the number of ROM banks for this ROM. Each bank is 16KB.
    pub fn rom_banks(&self) -> Option<uint> {
        let id = self.rom_byte(offsets::ROM_SIZE as u16);

        let nbanks =
            match id {
                0x00 => 2,
                0x01 => 4,
                0x02 => 8,
                0x03 => 16,
                0x04 => 32,
                0x05 => 64,
                0x06 => 128,
                0x52 => 72,
                0x53 => 80,
                0x54 => 96,
                // Unknown value
                _    => return None,
            };

        Some(nbanks)
    }

    /// Return the number of RAM banks for this ROM along with the
    /// size of each bank in bytes.
    pub fn ram_banks(&self) -> Option<(uint, uint)> {
        let id = self.rom_byte(offsets::RAM_SIZE as u16);

        let (nbanks, bank_size_kb) =
            match id {
                0x00 => (0,  0),
                0x01 => (1,  2),
                0x02 => (1,  8),
                0x03 => (4,  8),
                0x04 => (16, 8),
                // Unknown value
                _    => return None,
            };

        Some((nbanks, bank_size_kb * 1024))
    }

    pub fn rom_byte(&self, offset: u16) -> u8 {
        let off = offset as uint;

        if off < ROM_BANK_SIZE {
            self.rom[off]
        } else {
            self.rom[self.rom_offset as uint + off]
        }
    }

    pub fn set_rom_byte(&mut self, offset: u16, val: u8) {
        // Let specialized cartridge type handle that
        (self.model.write)(self, offset, val)
    }

    pub fn ram_byte(&self, offset: u16) -> u8 {
        *self.ram.get(offset as uint).unwrap_or(&0)
    }

    pub fn set_ram_byte(&mut self, offset: u16, val: u8) {
        let addr = self.ram_offset + offset as uint;

        if self.ram_wp {
            debug!("Attempt to write to cartridge RAM while protected");
            return;
        }

        if let Some(b) = self.ram.get_mut(addr) {
            *b = val;
        }
    }

    /// Retrieve current ROM bank number for the bankable range at
    /// [0x4000, 0x7fff]
    pub fn rom_bank(&self) -> u8 {
        self.rom_bank
    }

    /// Set new ROM bank number for the bankable range at
    /// [0x4000, 0x7fff]
    pub fn set_rom_bank(&mut self, bank: u8) {
        self.rom_bank = bank;

        // Recompute offset value to avoid doing it at each ROM read.
        self.rom_offset = ROM_BANK_SIZE *
            match self.rom_bank {
                /// We can't select bank 0, it defaults to 1
                0 => 0,
                n => (n - 1) as uint,
            };
    }

    /// Enable or disable RAM write protect
    pub fn set_ram_wp(&mut self, wp: bool) {
        self.ram_wp = wp
    }

    /// Return the value of the `bank_ram` flag
    pub fn bank_ram(&self) -> bool {
        self.bank_ram
    }

    /// Set the value of the `bank_ram` flag
    pub fn set_bank_ram(&mut self, v: bool) {
        self.bank_ram = v
    }

    /// Set new RAM bank number
    pub fn set_ram_bank(&mut self, bank: u8) {
        // Bankable RAM is always 8KB per bank
        self.ram_offset = bank as uint * 8 * 1024;
    }
}

impl Drop for Cartridge {
    fn drop(&mut self) {
        // Update save file when Cartridge is dropped
        if let Err(e) = self.save_ram() {
            // Display the error but don't panic since we might
            // already be in the middle of a panic unwinding
            println!("Couldn't save: {}", e);
        }
    }
}

impl Show for Cartridge {
    fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
        let name = match self.name() {
            Some(s) => s,
            None    => "<INVALID>".to_string(),
        };

        let rombanks = match self.rom_banks() {
            Some(n) => n,
            None    => 0,
        };

        let (rambanks, rambanksize) = match self.ram_banks() {
            Some(n) => n,
            None    => (0, 0),
        };

        try!(write!(f,
                    "'{}' (Model: {}, \
                           ROM banks: {}, \
                           RAM banks: {}, \
                           RAM bank size: {}KB)",
                    name, self.model.name, rombanks, rambanks, rambanksize));

        Ok(())
    }
}

// Each ROM bank is always 16KB
const ROM_BANK_SIZE: uint = 16 * 1024;

mod offsets {
    //! Various offset values to access special memory locations within the ROM

    /// Title. Upper case ASCII 16bytes long, padded with 0s if shorter
    pub const TITLE:    uint = 0x134;
    /// Cartridge type
    pub const TYPE:     uint = 0x147;
    pub const ROM_SIZE: uint = 0x148;
    pub const RAM_SIZE: uint = 0x149;
}
<commit_msg>Fix bogus address calculation when read cartridge RAM<commit_after>//! Cartridge emulation. There are multiple cartridge types with
//! different capabilities (bankable ROM/RAM, battery, RTC etc...).

use std::fmt::{Show, Formatter, Error};
use std::io::{File, Reader, Writer, IoResult, Open, ReadWrite, SeekSet};

mod models;

/// Common state for all cartridge types
pub struct Cartridge {
    /// Cartridge ROM data
    rom:        Vec<u8>,
    /// Cartridge RAM data
    ram:        Vec<u8>,
    /// Current number of the rom bank mapped at [0x4000, 0x7fff]
    rom_bank:   u8,
    /// Current bank offset for the bank mapped at [0x4000, 0x7fff].
    /// This value is added to ROM register addresses when they're in
    /// that range.
    rom_offset: uint,
    /// Current bank offset for the RAM
    ram_offset: uint,
    /// If `true` RAM is write protected
    ram_wp:     bool,
    /// Certain cartridges allow banking either the RAM or ROM
    /// depending on the value of this flag.
    bank_ram:   bool,
    /// struct used to handle model specific functions
    model:      models::Model,
    /// Path to the ROM image for this cartridge
    path:       Path,
    /// optional save file used to store non-volatile RAM on emulator
    /// shutdown
    save_file:  Option<File>,
}

impl Cartridge {
    /// Load a Cartridge ROM from `path`.
    pub fn from_path(rom_path: &Path) -> IoResult<Cartridge> {
        let mut source = try!(File::open(rom_path));

        // There must always be at least two ROM banks
        let rom = try!(source.read_exact(2 * ROM_BANK_SIZE));

        let model = models::from_id(rom[offsets::TYPE]);

        let mut cartridge = Cartridge {
            rom:        rom,
            ram:        Vec::new(),
            // Default to bank 1 for bankable region
            rom_bank:   1,
            rom_offset: 0,
            ram_offset: 0,
            ram_wp:     false,
            bank_ram:   false,
            model:      model,
            path:       rom_path.clone(),
            save_file:  None,
        };

        let rombanks = match cartridge.rom_banks() {
            Some(n) => n,
            None    => panic!("Can't determine ROM size"),
        };

        // Read the remaining roms banks
        if rombanks > 2 {
            let remb      = rombanks - 2;
            let mut off   = 2    * ROM_BANK_SIZE;
            let mut remsz = remb * ROM_BANK_SIZE;

            // Reserve space for the remaining banks
            cartridge.rom.grow(remsz, 0);

            while remsz > 0 {
                let r = try!(source.read(cartridge.rom.slice_from_mut(off)));

                remsz -= r;
                off   += r;
            }
        }

        try!(cartridge.init_ram());

        Ok(cartridge)
    }

    /// Init cartridge RAM and tie it with a `File` for saving if
    /// necessary.
    fn init_ram(&mut self) -> IoResult<()> {
        let (rambanks, banksize) = match self.ram_banks() {
            Some(v) => v,
            None    => panic!("Can't determine RAM size"),
        };

        let ramsize = rambanks * banksize;

        if ramsize == 0 {
            // No RAM on this cartridge, we're done
            return Ok(());
        }

        // We have some RAM, open the save file or create it if it
        // doesn't exist yet
        let mut savepath = self.path.clone();
        savepath.set_extension("sav");

        let mut save_file = try!(File::open_mode(&savepath,
                                                Open,
                                                ReadWrite));

        let save_size = try!(save_file.stat()).size;

        if save_size == 0 {
            // The file is empty (probably new). initialize
            // the RAM with 0s.
            self.ram.grow(ramsize, 0);
            // Then fill the file with the right amount of 0s
            // to reserve enough space for saving later.
            try!(save_file.write(self.ram.as_slice()));
        } else if save_size == (ramsize as u64) {
            // The file contains a RAM image
            self.ram = try!(save_file.read_exact(ramsize as uint));
        } else {
            panic!("Unexpected save file size for {}: expected {} got {}",
                   savepath.display(), ramsize, save_size);
        }

        // Store the file handle to save progress later
        self.save_file = Some(save_file);

        Ok(())
    }

    /// Update the save file
    pub fn save_ram(&mut self) -> IoResult<()> {
        if let Some(mut f) = self.save_file.as_mut() {
            // Rewind to the beginning of the file and update its
            // contents
            println!("Saving to {}", f.path().display());

            try!(f.seek(0, SeekSet));
            try!(f.write(self.ram.as_slice()));
        }

        Ok(())
    }

    /// Attempt to retreive the rom's name
    pub fn name(&self) -> Option<String> {
        let mut name = String::with_capacity(16);

        for i in range(0, 16) {
            let c = self.rom[offsets::TITLE + i].to_ascii();

            // If the name is shorter than 16bytes it's padded with 0s
            if c == 0.to_ascii() {
                break;
            }

            // Only uppercase ASCII is valid
            if !(c.is_uppercase() || c.is_blank() || c.is_punctuation()) {
                return None;
            }

            // Append new character
            name.grow(1, c.as_char());
        }

        Some(name)
    }

    /// Return the number of ROM banks for this ROM. Each bank is 16KB.
    pub fn rom_banks(&self) -> Option<uint> {
        let id = self.rom_byte(offsets::ROM_SIZE as u16);

        let nbanks =
            match id {
                0x00 => 2,
                0x01 => 4,
                0x02 => 8,
                0x03 => 16,
                0x04 => 32,
                0x05 => 64,
                0x06 => 128,
                0x52 => 72,
                0x53 => 80,
                0x54 => 96,
                // Unknown value
                _    => return None,
            };

        Some(nbanks)
    }

    /// Return the number of RAM banks for this ROM along with the
    /// size of each bank in bytes.
    pub fn ram_banks(&self) -> Option<(uint, uint)> {
        let id = self.rom_byte(offsets::RAM_SIZE as u16);

        let (nbanks, bank_size_kb) =
            match id {
                0x00 => (0,  0),
                0x01 => (1,  2),
                0x02 => (1,  8),
                0x03 => (4,  8),
                0x04 => (16, 8),
                // Unknown value
                _    => return None,
            };

        Some((nbanks, bank_size_kb * 1024))
    }

    pub fn rom_byte(&self, offset: u16) -> u8 {
        let off = offset as uint;

        if off < ROM_BANK_SIZE {
            self.rom[off]
        } else {
            self.rom[self.rom_offset as uint + off]
        }
    }

    pub fn set_rom_byte(&mut self, offset: u16, val: u8) {
        // Let specialized cartridge type handle that
        (self.model.write)(self, offset, val)
    }

    pub fn ram_byte(&self, offset: u16) -> u8 {
        let addr = self.ram_offset + offset as uint;

        *self.ram.get(addr).unwrap_or(&0)
    }

    pub fn set_ram_byte(&mut self, offset: u16, val: u8) {
        let addr = self.ram_offset + offset as uint;

        if self.ram_wp {
            debug!("Attempt to write to cartridge RAM while protected");
            return;
        }

        if let Some(b) = self.ram.get_mut(addr) {
            *b = val;
        }
    }

    /// Retrieve current ROM bank number for the bankable range at
    /// [0x4000, 0x7fff]
    pub fn rom_bank(&self) -> u8 {
        self.rom_bank
    }

    /// Set new ROM bank number for the bankable range at
    /// [0x4000, 0x7fff]
    pub fn set_rom_bank(&mut self, bank: u8) {
        self.rom_bank = bank;

        // Recompute offset value to avoid doing it at each ROM read.
        self.rom_offset = ROM_BANK_SIZE *
            match self.rom_bank {
                /// We can't select bank 0, it defaults to 1
                0 => 0,
                n => (n - 1) as uint,
            };
    }

    /// Enable or disable RAM write protect
    pub fn set_ram_wp(&mut self, wp: bool) {
        self.ram_wp = wp
    }

    /// Return the value of the `bank_ram` flag
    pub fn bank_ram(&self) -> bool {
        self.bank_ram
    }

    /// Set the value of the `bank_ram` flag
    pub fn set_bank_ram(&mut self, v: bool) {
        self.bank_ram = v
    }

    /// Set new RAM bank number
    pub fn set_ram_bank(&mut self, bank: u8) {
        // Bankable RAM is always 8KB per bank
        self.ram_offset = bank as uint * 8 * 1024;
    }
}

impl Drop for Cartridge {
    fn drop(&mut self) {
        // Update save file when Cartridge is dropped
        if let Err(e) = self.save_ram() {
            // Display the error but don't panic since we might
            // already be in the middle of a panic unwinding
            println!("Couldn't save: {}", e);
        }
    }
}

impl Show for Cartridge {
    fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
        let name = match self.name() {
            Some(s) => s,
            None    => "<INVALID>".to_string(),
        };

        let rombanks = match self.rom_banks() {
            Some(n) => n,
            None    => 0,
        };

        let (rambanks, rambanksize) = match self.ram_banks() {
            Some(n) => n,
            None    => (0, 0),
        };

        try!(write!(f,
                    "'{}' (Model: {}, \
                           ROM banks: {}, \
                           RAM banks: {}, \
                           RAM bank size: {}KB)",
                    name, self.model.name, rombanks, rambanks, rambanksize));

        Ok(())
    }
}

// Each ROM bank is always 16KB
const ROM_BANK_SIZE: uint = 16 * 1024;

mod offsets {
    //! Various offset values to access special memory locations within the ROM

    /// Title. Upper case ASCII 16bytes long, padded with 0s if shorter
    pub const TITLE:    uint = 0x134;
    /// Cartridge type
    pub const TYPE:     uint = 0x147;
    pub const ROM_SIZE: uint = 0x148;
    pub const RAM_SIZE: uint = 0x149;
}
<|endoftext|>


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

use crate::{db::MultiThreaded, ffi, Options};

use std::sync::Arc;

/// The name of the default column family.
///
/// The column family with this name is created implicitly whenever column
/// families are used.
pub const DEFAULT_COLUMN_FAMILY_NAME: &str = "default";

/// A descriptor for a RocksDB column family.
///
/// A description of the column family, containing the name and `Options`.
pub struct ColumnFamilyDescriptor {
    pub(crate) name: String,
    pub(crate) options: Options,
}

impl ColumnFamilyDescriptor {
    // Create a new column family descriptor with the specified name and options.
    pub fn new<S>(name: S, options: Options) -> Self
    where
        S: Into<String>,
    {
        Self {
            name: name.into(),
            options,
        }
    }
}

/// An opaque type used to represent a column family. Returned from some functions, and used
/// in others
pub struct ColumnFamily {
    pub(crate) inner: *mut ffi::rocksdb_column_family_handle_t,
}

/// A specialized opaque type used to represent a column family by the [`MultiThreaded`]
/// mode. Clone (and Copy) is derived to behave like `&ColumnFamily` (this is used for
/// single-threaded mode). `Clone`/`Copy` is safe because this lifetime is bound to DB like
/// iterators/snapshots. On top of it, this is as cheap and small as `&ColumnFamily` because
/// this only has a single pointer-wide field.
pub struct BoundColumnFamily<'a> {
    pub(crate) inner: *mut ffi::rocksdb_column_family_handle_t,
    pub(crate) multi_threaded_cfs: std::marker::PhantomData<&'a MultiThreaded>,
}

// internal struct which isn't exposed to public api.
// but its memory will be exposed after transmute()-ing to BoundColumnFamily.
// ColumnFamily's lifetime should be bound to DB. But, db holds cfs and cfs can't easily
// self-reference DB as its lifetime due to rust's type system
pub(crate) struct UnboundColumnFamily {
    pub(crate) inner: *mut ffi::rocksdb_column_family_handle_t,
}

impl UnboundColumnFamily {
    pub(crate) fn bound_column_family<'a>(self: Arc<Self>) -> Arc<BoundColumnFamily<'a>> {
        // SAFETY: the new BoundColumnFamily here just adding lifetime,
        // so that column family handle won't outlive db.
        unsafe { std::mem::transmute(self) }
    }
}

fn destroy_handle(handle: *mut ffi::rocksdb_column_family_handle_t) {
    // SAFETY: This should be called only from various Drop::drop(), strictly keeping a 1-to-1
    // ownership to avoid double invocation to the rocksdb function with same handle.
    unsafe {
        ffi::rocksdb_column_family_handle_destroy(handle);
    }
}

impl Drop for ColumnFamily {
    fn drop(&mut self) {
        destroy_handle(self.inner);
    }
}

// these behaviors must be identical between BoundColumnFamily and UnboundColumnFamily
// due to the unsafe transmute() in bound_column_family()!
impl<'a> Drop for BoundColumnFamily<'a> {
    fn drop(&mut self) {
        destroy_handle(self.inner);
    }
}

impl Drop for UnboundColumnFamily {
    fn drop(&mut self) {
        destroy_handle(self.inner);
    }
}

/// Handy type alias to hide actual type difference to reference [`ColumnFamily`]
/// depending on the `multi-threaded-cf` crate feature.
#[cfg(not(feature = "multi-threaded-cf"))]
pub type ColumnFamilyRef<'a> = &'a ColumnFamily;

#[cfg(feature = "multi-threaded-cf")]
pub type ColumnFamilyRef<'a> = Arc<BoundColumnFamily<'a>>;

/// Utility trait to accept both supported references to `ColumnFamily`
/// (`&ColumnFamily` and `BoundColumnFamily`)
pub trait AsColumnFamilyRef {
    fn inner(&self) -> *mut ffi::rocksdb_column_family_handle_t;
}

impl AsColumnFamilyRef for ColumnFamily {
    fn inner(&self) -> *mut ffi::rocksdb_column_family_handle_t {
        self.inner
    }
}

impl<'a> AsColumnFamilyRef for &'a ColumnFamily {
    fn inner(&self) -> *mut ffi::rocksdb_column_family_handle_t {
        self.inner
    }
}

// Only implement for Arc-ed BoundColumnFamily as this tightly coupled and
// implementation detail, considering use of std::mem::transmute. BoundColumnFamily
// isn't expected to be used as naked.
// Also, ColumnFamilyRef might not be Arc<BoundColumnFamily<'a>> depending crate
// feature flags so, we can't use the type alias here.
impl<'a> AsColumnFamilyRef for Arc<BoundColumnFamily<'a>> {
    fn inner(&self) -> *mut ffi::rocksdb_column_family_handle_t {
        self.inner
    }
}

unsafe impl Send for ColumnFamily {}
unsafe impl Send for UnboundColumnFamily {}
unsafe impl Sync for UnboundColumnFamily {}
unsafe impl<'a> Send for BoundColumnFamily<'a> {}
<commit_msg>Fix theoretical UB while transmuting Arc (#588)<commit_after>// Copyright 2020 Tyler Neely
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::{db::MultiThreaded, ffi, Options};

use std::sync::Arc;

/// The name of the default column family.
///
/// The column family with this name is created implicitly whenever column
/// families are used.
pub const DEFAULT_COLUMN_FAMILY_NAME: &str = "default";

/// A descriptor for a RocksDB column family.
///
/// A description of the column family, containing the name and `Options`.
pub struct ColumnFamilyDescriptor {
    pub(crate) name: String,
    pub(crate) options: Options,
}

impl ColumnFamilyDescriptor {
    // Create a new column family descriptor with the specified name and options.
    pub fn new<S>(name: S, options: Options) -> Self
    where
        S: Into<String>,
    {
        Self {
            name: name.into(),
            options,
        }
    }
}

/// An opaque type used to represent a column family. Returned from some functions, and used
/// in others
pub struct ColumnFamily {
    pub(crate) inner: *mut ffi::rocksdb_column_family_handle_t,
}

/// A specialized opaque type used to represent a column family by the [`MultiThreaded`]
/// mode. Clone (and Copy) is derived to behave like `&ColumnFamily` (this is used for
/// single-threaded mode). `Clone`/`Copy` is safe because this lifetime is bound to DB like
/// iterators/snapshots. On top of it, this is as cheap and small as `&ColumnFamily` because
/// this only has a single pointer-wide field.
pub struct BoundColumnFamily<'a> {
    pub(crate) inner: *mut ffi::rocksdb_column_family_handle_t,
    pub(crate) multi_threaded_cfs: std::marker::PhantomData<&'a MultiThreaded>,
}

// internal struct which isn't exposed to public api.
// but its memory will be exposed after transmute()-ing to BoundColumnFamily.
// ColumnFamily's lifetime should be bound to DB. But, db holds cfs and cfs can't easily
// self-reference DB as its lifetime due to rust's type system
pub(crate) struct UnboundColumnFamily {
    pub(crate) inner: *mut ffi::rocksdb_column_family_handle_t,
}

impl UnboundColumnFamily {
    pub(crate) fn bound_column_family<'a>(self: Arc<Self>) -> Arc<BoundColumnFamily<'a>> {
        // SAFETY: the new BoundColumnFamily here just adding lifetime,
        // so that column family handle won't outlive db.
        unsafe { Arc::from_raw(Arc::into_raw(self).cast()) }
    }
}

fn destroy_handle(handle: *mut ffi::rocksdb_column_family_handle_t) {
    // SAFETY: This should be called only from various Drop::drop(), strictly keeping a 1-to-1
    // ownership to avoid double invocation to the rocksdb function with same handle.
    unsafe {
        ffi::rocksdb_column_family_handle_destroy(handle);
    }
}

impl Drop for ColumnFamily {
    fn drop(&mut self) {
        destroy_handle(self.inner);
    }
}

// these behaviors must be identical between BoundColumnFamily and UnboundColumnFamily
// due to the unsafe transmute() in bound_column_family()!
impl<'a> Drop for BoundColumnFamily<'a> {
    fn drop(&mut self) {
        destroy_handle(self.inner);
    }
}

impl Drop for UnboundColumnFamily {
    fn drop(&mut self) {
        destroy_handle(self.inner);
    }
}

/// Handy type alias to hide actual type difference to reference [`ColumnFamily`]
/// depending on the `multi-threaded-cf` crate feature.
#[cfg(not(feature = "multi-threaded-cf"))]
pub type ColumnFamilyRef<'a> = &'a ColumnFamily;

#[cfg(feature = "multi-threaded-cf")]
pub type ColumnFamilyRef<'a> = Arc<BoundColumnFamily<'a>>;

/// Utility trait to accept both supported references to `ColumnFamily`
/// (`&ColumnFamily` and `BoundColumnFamily`)
pub trait AsColumnFamilyRef {
    fn inner(&self) -> *mut ffi::rocksdb_column_family_handle_t;
}

impl AsColumnFamilyRef for ColumnFamily {
    fn inner(&self) -> *mut ffi::rocksdb_column_family_handle_t {
        self.inner
    }
}

impl<'a> AsColumnFamilyRef for &'a ColumnFamily {
    fn inner(&self) -> *mut ffi::rocksdb_column_family_handle_t {
        self.inner
    }
}

// Only implement for Arc-ed BoundColumnFamily as this tightly coupled and
// implementation detail, considering use of std::mem::transmute. BoundColumnFamily
// isn't expected to be used as naked.
// Also, ColumnFamilyRef might not be Arc<BoundColumnFamily<'a>> depending crate
// feature flags so, we can't use the type alias here.
impl<'a> AsColumnFamilyRef for Arc<BoundColumnFamily<'a>> {
    fn inner(&self) -> *mut ffi::rocksdb_column_family_handle_t {
        self.inner
    }
}

unsafe impl Send for ColumnFamily {}
unsafe impl Send for UnboundColumnFamily {}
unsafe impl Sync for UnboundColumnFamily {}
unsafe impl<'a> Send for BoundColumnFamily<'a> {}
<|endoftext|>


<commit_before>use cranelift_module::*;
use crate::prelude::*;
use rustc::mir::interpret::{read_target_uint, AllocId, ConstValue, GlobalId};
use rustc::ty::Const;
use rustc_mir::interpret::{CompileTimeEvaluator, Memory};

#[derive(Default)]
pub struct ConstantCx {
    constants: HashMap<AllocId, DataId>,
    done: HashSet<DataId>,
}

impl ConstantCx {
    pub fn finalize<B: Backend>(mut self, module: &mut Module<B>) {
        for data_id in self.done.drain() {
            module.finalize_data(data_id);
        }
    }
}

pub fn codegen_static<'a, 'tcx: 'a, B: Backend>(_cx: &mut CodegenCx<'a, 'tcx, B>, def_id: DefId) {
    unimpl!("Unimplemented static mono item {:?}", def_id);
}

pub fn codegen_static_ref<'a, 'tcx: 'a>(
    _fx: &mut FunctionCx<'a, 'tcx>,
    static_: &Static<'tcx>,
) -> CPlace<'tcx> {
    unimpl!("Unimplemented static mono item {:?}", static_.def_id);
}

pub fn trans_promoted<'a, 'tcx: 'a>(
    fx: &mut FunctionCx<'a, 'tcx>,
    promoted: Promoted,
) -> CPlace<'tcx> {
    let const_ = fx
        .tcx
        .const_eval(ParamEnv::reveal_all().and(GlobalId {
            instance: fx.instance,
            promoted: Some(promoted),
        })).unwrap();

    let const_ = force_eval_const(fx, const_);
    trans_const_place(fx, const_)
}

pub fn trans_constant<'a, 'tcx: 'a>(
    fx: &mut FunctionCx<'a, 'tcx>,
    constant: &Constant<'tcx>,
) -> CValue<'tcx> {
    let const_ = fx.monomorphize(&constant.literal);
    let const_ = force_eval_const(fx, const_);
    trans_const_value(fx, const_)
}

fn force_eval_const<'a, 'tcx: 'a>(
    fx: &FunctionCx<'a, 'tcx>,
    const_: &'tcx Const<'tcx>,
) -> &'tcx Const<'tcx> {
    match const_.val {
        ConstValue::Unevaluated(def_id, ref substs) => {
            let param_env = ParamEnv::reveal_all();
            let instance = Instance::resolve(fx.tcx, param_env, def_id, substs).unwrap();
            let cid = GlobalId {
                instance,
                promoted: None,
            };
            fx.tcx.const_eval(param_env.and(cid)).unwrap()
        }
        _ => const_,
    }
}

fn trans_const_value<'a, 'tcx: 'a>(
    fx: &mut FunctionCx<'a, 'tcx>,
    const_: &'tcx Const<'tcx>,
) -> CValue<'tcx> {
    let ty = fx.monomorphize(&const_.ty);
    let layout = fx.layout_of(ty);
    match ty.sty {
        TypeVariants::TyBool => {
            let bits = const_.val.to_bits(layout.size).unwrap();
            CValue::const_val(fx, ty, bits as u64 as i64)
        }
        TypeVariants::TyUint(_) => {
            let bits = const_.val.to_bits(layout.size).unwrap();
            CValue::const_val(fx, ty, bits as u64 as i64)
        }
        TypeVariants::TyInt(_) => {
            let bits = const_.val.to_bits(layout.size).unwrap();
            CValue::const_val(fx, ty, bits as i128 as i64)
        }
        TypeVariants::TyFnDef(def_id, substs) => {
            let func_ref = fx.get_function_ref(Instance::new(def_id, substs));
            CValue::Func(func_ref, layout)
        }
        _ => trans_const_place(fx, const_).to_cvalue(fx),
    }
}

fn trans_const_place<'a, 'tcx: 'a>(
    fx: &mut FunctionCx<'a, 'tcx>,
    const_: &'tcx Const<'tcx>,
) -> CPlace<'tcx> {
    let ty = fx.monomorphize(&const_.ty);
    let layout = fx.layout_of(ty);
    if !fx
        .tcx
        .sess
        .crate_types
        .get()
        .contains(&CrateType::Executable)
    {
        // TODO: cranelift-module api seems to be used wrong,
        // thus causing panics for some consts, so this disables it
        return CPlace::Addr(fx.bcx.ins().iconst(types::I64, 0), layout);
    }
    let alloc = fx.tcx.const_value_to_allocation(const_);
    //println!("const value: {:?} allocation: {:?}", value, alloc);
    let alloc_id = fx.tcx.alloc_map.lock().allocate(alloc);
    let data_id = get_global_for_alloc_id(
        fx.tcx,
        fx.module,
        fx.constants,
        alloc_id,
    );
    let local_data_id = fx.module.declare_data_in_func(data_id, &mut fx.bcx.func);
    // TODO: does global_value return a ptr of a val?
    let global_ptr = fx.bcx.ins().global_value(types::I64, local_data_id);
    CPlace::Addr(global_ptr, layout)
}

// If ret.1 is true, then the global didn't exist before
fn define_global_for_alloc_id<'a, 'tcx: 'a, B: Backend>(
    module: &mut Module<B>,
    cx: &mut ConstantCx,//<'a, 'tcx>,
    alloc_id: AllocId,
    todo: &mut HashMap<AllocId, DataId>,
) -> (DataId, bool) {
    use std::collections::hash_map::Entry;
    match cx.constants.entry(alloc_id) {
        Entry::Occupied(mut occ) => (*occ.get_mut(), false),
        Entry::Vacant(vac) => {
            let data_id = module
                .declare_data(&alloc_id.0.to_string(), Linkage::Local, false)
                .unwrap();
            todo.insert(alloc_id, data_id);
            vac.insert(data_id);
            (data_id, true)
        }
    }
}

fn get_global_for_alloc_id<'a, 'tcx: 'a, B: Backend + 'a>(
    tcx: TyCtxt<'a, 'tcx, 'tcx>,
    module: &mut Module<B>,
    cx: &mut ConstantCx,//<'a, 'tcx>,
    alloc_id: AllocId,
) -> DataId {
    if let Some(data_id) = cx.constants.get(&alloc_id) {
        return *data_id;
    }

    let memory = Memory::<CompileTimeEvaluator>::new(tcx.at(DUMMY_SP), ());

    let mut todo = HashMap::new();
    define_global_for_alloc_id(module, cx, alloc_id, &mut todo);

    while let Some((alloc_id, data_id)) = {
        let next = todo.drain().next();
        next
    } {
        println!(
            "cur: {:?}:{:?} todo: {:?} done: {:?}",
            alloc_id, data_id, todo, cx.done
        );

        let alloc = memory.get(alloc_id).unwrap();
        //let alloc = tcx.alloc_map.lock().get(alloc_id).unwrap();
        let mut data_ctx = DataContext::new();

        data_ctx.define(
            alloc.bytes.to_vec().into_boxed_slice(),
            Writability::Readonly,
        );

        for &(offset, reloc) in alloc.relocations.iter() {
            let data_id = define_global_for_alloc_id(module, cx, reloc, &mut todo).0;

            let reloc_offset = {
                let endianness = memory.endianness();
                let offset = offset.bytes() as usize;
                let ptr_size = tcx.data_layout.pointer_size;
                let bytes = &alloc.bytes[offset..offset + ptr_size.bytes() as usize];
                read_target_uint(endianness, bytes).unwrap()
            };

            // TODO: is this a correct usage of the api
            let global_value = module.declare_data_in_data(data_id, &mut data_ctx);
            data_ctx.write_data_addr(reloc_offset as u32, global_value, 0);
        }

        module.define_data(data_id, &data_ctx).unwrap();
        cx.done.insert(data_id);
    }
    *cx.constants.get(&alloc_id).unwrap()
}
<commit_msg>Simplify define_global_for_alloc_id<commit_after>use cranelift_module::*;
use crate::prelude::*;
use rustc::mir::interpret::{read_target_uint, AllocId, ConstValue, GlobalId};
use rustc::ty::Const;
use rustc_mir::interpret::{CompileTimeEvaluator, Memory};

#[derive(Default)]
pub struct ConstantCx {
    constants: HashMap<AllocId, DataId>,
    done: HashSet<DataId>,
}

impl ConstantCx {
    pub fn finalize<B: Backend>(mut self, module: &mut Module<B>) {
        for data_id in self.done.drain() {
            module.finalize_data(data_id);
        }
    }
}

pub fn codegen_static<'a, 'tcx: 'a, B: Backend>(_cx: &mut CodegenCx<'a, 'tcx, B>, def_id: DefId) {
    unimpl!("Unimplemented static mono item {:?}", def_id);
}

pub fn codegen_static_ref<'a, 'tcx: 'a>(
    _fx: &mut FunctionCx<'a, 'tcx>,
    static_: &Static<'tcx>,
) -> CPlace<'tcx> {
    unimpl!("Unimplemented static mono item {:?}", static_.def_id);
}

pub fn trans_promoted<'a, 'tcx: 'a>(
    fx: &mut FunctionCx<'a, 'tcx>,
    promoted: Promoted,
) -> CPlace<'tcx> {
    let const_ = fx
        .tcx
        .const_eval(ParamEnv::reveal_all().and(GlobalId {
            instance: fx.instance,
            promoted: Some(promoted),
        })).unwrap();

    let const_ = force_eval_const(fx, const_);
    trans_const_place(fx, const_)
}

pub fn trans_constant<'a, 'tcx: 'a>(
    fx: &mut FunctionCx<'a, 'tcx>,
    constant: &Constant<'tcx>,
) -> CValue<'tcx> {
    let const_ = fx.monomorphize(&constant.literal);
    let const_ = force_eval_const(fx, const_);
    trans_const_value(fx, const_)
}

fn force_eval_const<'a, 'tcx: 'a>(
    fx: &FunctionCx<'a, 'tcx>,
    const_: &'tcx Const<'tcx>,
) -> &'tcx Const<'tcx> {
    match const_.val {
        ConstValue::Unevaluated(def_id, ref substs) => {
            let param_env = ParamEnv::reveal_all();
            let instance = Instance::resolve(fx.tcx, param_env, def_id, substs).unwrap();
            let cid = GlobalId {
                instance,
                promoted: None,
            };
            fx.tcx.const_eval(param_env.and(cid)).unwrap()
        }
        _ => const_,
    }
}

fn trans_const_value<'a, 'tcx: 'a>(
    fx: &mut FunctionCx<'a, 'tcx>,
    const_: &'tcx Const<'tcx>,
) -> CValue<'tcx> {
    let ty = fx.monomorphize(&const_.ty);
    let layout = fx.layout_of(ty);
    match ty.sty {
        TypeVariants::TyBool => {
            let bits = const_.val.to_bits(layout.size).unwrap();
            CValue::const_val(fx, ty, bits as u64 as i64)
        }
        TypeVariants::TyUint(_) => {
            let bits = const_.val.to_bits(layout.size).unwrap();
            CValue::const_val(fx, ty, bits as u64 as i64)
        }
        TypeVariants::TyInt(_) => {
            let bits = const_.val.to_bits(layout.size).unwrap();
            CValue::const_val(fx, ty, bits as i128 as i64)
        }
        TypeVariants::TyFnDef(def_id, substs) => {
            let func_ref = fx.get_function_ref(Instance::new(def_id, substs));
            CValue::Func(func_ref, layout)
        }
        _ => trans_const_place(fx, const_).to_cvalue(fx),
    }
}

fn trans_const_place<'a, 'tcx: 'a>(
    fx: &mut FunctionCx<'a, 'tcx>,
    const_: &'tcx Const<'tcx>,
) -> CPlace<'tcx> {
    let ty = fx.monomorphize(&const_.ty);
    let layout = fx.layout_of(ty);
    if !fx
        .tcx
        .sess
        .crate_types
        .get()
        .contains(&CrateType::Executable)
    {
        // TODO: cranelift-module api seems to be used wrong,
        // thus causing panics for some consts, so this disables it
        return CPlace::Addr(fx.bcx.ins().iconst(types::I64, 0), layout);
    }
    let alloc = fx.tcx.const_value_to_allocation(const_);
    //println!("const value: {:?} allocation: {:?}", value, alloc);
    let alloc_id = fx.tcx.alloc_map.lock().allocate(alloc);
    let data_id = get_global_for_alloc_id(
        fx.tcx,
        fx.module,
        fx.constants,
        alloc_id,
    );
    let local_data_id = fx.module.declare_data_in_func(data_id, &mut fx.bcx.func);
    // TODO: does global_value return a ptr of a val?
    let global_ptr = fx.bcx.ins().global_value(types::I64, local_data_id);
    CPlace::Addr(global_ptr, layout)
}

// If ret.1 is true, then the global didn't exist before
fn define_global_for_alloc_id<'a, 'tcx: 'a, B: Backend>(
    module: &mut Module<B>,
    cx: &mut ConstantCx,
    alloc_id: AllocId,
    todo: &mut HashMap<AllocId, DataId>,
) -> DataId {
    *cx.constants.entry(alloc_id).or_insert_with(|| {
        let data_id = module
            .declare_data(&alloc_id.0.to_string(), Linkage::Local, false)
            .unwrap();
        todo.insert(alloc_id, data_id);
        data_id
    })
}

fn get_global_for_alloc_id<'a, 'tcx: 'a, B: Backend + 'a>(
    tcx: TyCtxt<'a, 'tcx, 'tcx>,
    module: &mut Module<B>,
    cx: &mut ConstantCx,//<'a, 'tcx>,
    alloc_id: AllocId,
) -> DataId {
    if let Some(data_id) = cx.constants.get(&alloc_id) {
        return *data_id;
    }

    let memory = Memory::<CompileTimeEvaluator>::new(tcx.at(DUMMY_SP), ());

    let mut todo = HashMap::new();
    define_global_for_alloc_id(module, cx, alloc_id, &mut todo);

    while let Some((alloc_id, data_id)) = {
        let next = todo.drain().next();
        next
    } {
        println!(
            "cur: {:?}:{:?} todo: {:?} done: {:?}",
            alloc_id, data_id, todo, cx.done
        );

        let alloc = memory.get(alloc_id).unwrap();
        //let alloc = tcx.alloc_map.lock().get(alloc_id).unwrap();
        let mut data_ctx = DataContext::new();

        data_ctx.define(
            alloc.bytes.to_vec().into_boxed_slice(),
            Writability::Readonly,
        );

        for &(offset, reloc) in alloc.relocations.iter() {
            let data_id = define_global_for_alloc_id(module, cx, reloc, &mut todo);

            let reloc_offset = {
                let endianness = memory.endianness();
                let offset = offset.bytes() as usize;
                let ptr_size = tcx.data_layout.pointer_size;
                let bytes = &alloc.bytes[offset..offset + ptr_size.bytes() as usize];
                read_target_uint(endianness, bytes).unwrap()
            };

            // TODO: is this a correct usage of the api
            let global_value = module.declare_data_in_data(data_id, &mut data_ctx);
            data_ctx.write_data_addr(reloc_offset as u32, global_value, 0);
        }

        module.define_data(data_id, &data_ctx).unwrap();
        cx.done.insert(data_id);
    }
    *cx.constants.get(&alloc_id).unwrap()
}
<|endoftext|>


<commit_before>extern crate libc;
use self::libc::*;

use low_level as ll;
use error::Error;
use device::Device;
use context::Context;
use event::{Event, EventStatus, EventLike};
use kernel::Kernel;
use mem::Mem;

use std::ptr;
use std::mem::{size_of, transmute};

pub struct CommandQueue {
    pub id: ll::CommandQueue,
}

pub struct ReadLock<'a, T: Sized + 'a> {
    #[allow(dead_code)]
    sl: &'a [T],
    evt: Event,
}

pub struct WriteLock<'a, T: Sized + 'a> {
    #[allow(dead_code)]
    sl: &'a mut [T],
    evt: Event,
}

impl<'a, T: Sized> Drop for ReadLock<'a, T> {
    fn drop(self: &mut Self) -> () {
        self.evt.wait().ok().expect(
            "error waiting to finish write for ReadLock");
    }
}

impl<'a, T: Sized> EventLike for ReadLock<'a, T> {
    fn status(self: &Self) -> Result<EventStatus, Error> {
        self.evt.status()
    }

    fn wait(self: &Self) -> Result<(), Error> {
        self.evt.wait()
    }
}

impl<'a, T: Sized> Drop for WriteLock<'a, T> {
    fn drop(self: &mut Self) -> () {
        self.evt.wait().ok().expect(
            "error waiting to finish write for ReadLock");
    }
}

impl<'a, T: Sized> EventLike for WriteLock<'a, T> {
    fn status(self: &Self) -> Result<EventStatus, Error> {
        self.evt.status()
    }

    fn wait(self: &Self) -> Result<(), Error> {
        self.evt.wait()
    }
}

unsafe impl Send for CommandQueue { }

impl CommandQueue {
    pub fn new(context: Context,
               device: Device) -> Result<CommandQueue, Error> {
        let props: c_ulong = if try!(device.out_of_order_supported()) {
            1
        } else {
            0
        };
        let id = unsafe {
            let mut err: i32 = 0;
            let id = ll::clCreateCommandQueue(context.id,
                                              device.id,
                                              props,
                                              &mut err);
            try!(Error::check(err));
            id
        };
        Ok(CommandQueue{ id: id })
    }

    fn get_info(self: &Self, param: u32) -> Result<*mut c_void, Error> {
        let mut tr: *mut c_void = ptr::null_mut();
        unsafe {
            try!(Error::check(ll::clGetCommandQueueInfo(self.id,
                                                   param,
                                                   size_of::<*mut c_void>() as size_t,
                                                   transmute(&mut tr),
                                                   ptr::null_mut())));
        }
        Ok(tr)
    }

    pub fn context(self: &Self) -> Result<Context, Error> {
        let val = try!(self.get_info(0x1090));
        let id = unsafe { transmute(val) };
        unsafe {
            try!(Error::check(ll::clRetainContext(id)));
        }
        Ok(Context{id: id})
    }

    pub fn device(self: &Self) -> Result<Device, Error> {
        let val = try!(self.get_info(0x1091));
        let id = unsafe { transmute(val) };
        Ok(Device{id: id})
    }

    fn alloc_ll_buffer(self: &Self,
                    flags: c_ulong,
                    size_bytes: usize) -> Result<ll::Mem, Error> {
        let context = try!(self.context());
        let id = unsafe {
            let mut err: i32 = 0;
            let id = ll::clCreateBuffer(context.id,
                                        flags,
                                        size_bytes as size_t,
                                        ptr::null_mut(),
                                        &mut err);
            try!(Error::check(err));
            id
        };
        Ok(id)
    }

    fn write_buffer_raw<T: Sized>(self: &Self,
                       buf: &mut Mem,
                       offset: usize,
                       size: usize,
                       data: &[T]) -> Result<Event, Error> {
        let mut event_id: ll::Event = ptr::null_mut();
        let host_ptr = &data[0] as *const T;

        unsafe {
            try!(Error::check(ll::clEnqueueWriteBuffer(self.id,
                                                       buf.id,
                                                       0, // don't block
                                                       offset as size_t,
                                                       size as size_t,
                                                       transmute(host_ptr),
                                                       0,
                                                       ptr::null(),
                                                       &mut event_id)));
        }

        let event = Event{id: event_id};
        Ok(event)
    }

    fn read_buffer_raw<T: Sized>(self: &Self,
                                 buf: &Mem,
                                 offset: usize,
                                 size: usize,
                                 data: &mut [T]) -> Result<Event, Error> {
        let mut event_id: ll::Event = ptr::null_mut();
        let host_ptr = &mut data[0] as *mut T;

        unsafe {
            try!(Error::check(ll::clEnqueueReadBuffer(self.id,
                                                      buf.id,
                                                      0, // don't block,
                                                      offset as size_t,
                                                      size as size_t,
                                                      transmute(host_ptr),
                                                      0,
                                                      ptr::null(),
                                                      &mut event_id)));
        }

        let event = Event{id: event_id};
        Ok(event)
    }

    pub fn create_buffer(self: &Self,
                               size_bytes: usize) -> Result<Mem, Error> {
        let id = try!(self.alloc_ll_buffer(0, size_bytes));
        let tr = Mem::new(id);
        Ok(tr)
    }

    pub fn create_buffer_from_slice<T: Sized>(self: &Self,
                                    slice: &[T]) -> Result<Mem, Error> {
        let num_bytes = slice.len() * size_of::<T>();
        let mut buf = try!(self.create_buffer(num_bytes));
        try!(self.write_buffer(&mut buf, slice));
        Ok(buf)
    }

    pub fn write_buffer_unscoped<T: Sized>(self: &Self,
                                  mem: &mut Mem,
                                  slice: &[T]) -> Result<Event, Error> {
        let size_bytes = size_of::<T>()*slice.len();
        let evt = try!(self.write_buffer_raw(mem,
                                             0,
                                             size_bytes,
                                             slice));
        Ok(evt)
    }

    pub fn read_buffer_unscoped<T: Sized>(self: &Self,
                                                 mem: &Mem,
                                                 slice: &mut [T]) -> Result<Event, Error> {
        let size_bytes = size_of::<T>()*slice.len();
        let evt = try!(self.read_buffer_raw(mem,
                                            0,
                                            size_bytes,
                                            slice));
        Ok(evt)
    }

    pub fn write_buffer<'a, 'b, T: Sized + 'a>(self: &Self,
                                  mem: &'b mut Mem,
                                  slice: &'a [T]) -> Result<ReadLock<'a, T>, Error> {
        let size_bytes = size_of::<T>()*slice.len();
        let evt = try!(self.write_buffer_raw(mem,
                                             0,
                                             size_bytes,
                                             slice));
        Ok(ReadLock{ evt: evt, sl: slice })
    }

    pub fn read_buffer<'a, T: Sized>(self: &Self,
                                 mem: &Mem,
                                 slice: &'a mut [T]) -> Result<WriteLock<'a, T>, Error> {
        let size_bytes = size_of::<T>()*slice.len();
        let evt = try!(self.read_buffer_raw(mem,
                                            0,
                                            size_bytes,
                                            slice));
        Ok(WriteLock{ evt: evt, sl: slice })
    }

    fn next_mul(self: &Self, global: usize, local: usize) -> usize {
        if global % local == 0 {
            global
        } else {
            ((global/local)+1)*local
        }
    }

    pub fn run(
                self: &Self,
                kernel: &mut Kernel,
                local_size: (usize, usize, usize),
                global_size: (usize, usize, usize)) -> Result<Event, Error> {
        let global_size_fixed = (self.next_mul(global_size.0, local_size.0),
                                self.next_mul(global_size.1, local_size.1),
                                self.next_mul(global_size.2, local_size.2));

        let local_vec = vec!(local_size.0 as size_t,
                             local_size.1 as size_t,
                             local_size.2 as size_t);
        let global_vec = vec!(global_size_fixed.0 as size_t,
                              global_size_fixed.1 as size_t,
                              global_size_fixed.2 as size_t);

        let mut event_id: ll::Event = ptr::null_mut();
        unsafe {
            let events_val = ptr::null_mut();
            try!(Error::check(ll::clEnqueueNDRangeKernel(self.id,
                                                         kernel.id,
                                                         3,
                                                         ptr::null(),
                                                         &global_vec[0],
                                                         &local_vec[0],
                                                         0,
                                                         events_val,
                                                         &mut event_id)));
        }
        let event = Event{id: event_id};

        Ok(event)
    }

    pub fn run_with_events(
                self: &Self,
                kernel: &mut Kernel,
                local_size: (usize, usize, usize),
                global_size: (usize, usize, usize),
                events_in: &[Event]) -> Result<Event, Error> {
        let events = events_in.to_owned();
        let global_size_fixed = (self.next_mul(global_size.0, local_size.0),
                                self.next_mul(global_size.1, local_size.1),
                                self.next_mul(global_size.2, local_size.2));

        let local_vec = vec!(local_size.0 as size_t,
                             local_size.1 as size_t,
                             local_size.2 as size_t);
        let global_vec = vec!(global_size_fixed.0 as size_t,
                              global_size_fixed.1 as size_t,
                              global_size_fixed.2 as size_t);

        let mut event_id: ll::Event = ptr::null_mut();
        unsafe {
            let events_val = if events.len() > 0 {
                transmute(&events[0])
            } else {
                ptr::null_mut()
            };
            try!(Error::check(ll::clEnqueueNDRangeKernel(self.id,
                                                         kernel.id,
                                                         3,
                                                         ptr::null(),
                                                         &global_vec[0],
                                                         &local_vec[0],
                                                         events.len() as u32,
                                                         events_val,
                                                         &mut event_id)));
        }
        let event = Event{id: event_id};

        Ok(event)
    }
}

impl Drop for CommandQueue {
    fn drop(self: &mut Self) -> () {
        unsafe {
            Error::check(ll::clReleaseCommandQueue(self.id))
                .ok().expect("Error releasing command queue");
        }
    }
}

impl Clone for CommandQueue {
    fn clone(self: &Self) -> Self {
        unsafe {
            Error::check(ll::clRetainCommandQueue(self.id))
                .ok().expect("Error retaining command queue");
        }
        CommandQueue{ id: self.id }
    }
}

<commit_msg>another attempt to squash queued kernel bug<commit_after>extern crate libc;
use self::libc::*;

use low_level as ll;
use error::Error;
use device::Device;
use context::Context;
use event::{Event, EventStatus, EventLike};
use kernel::Kernel;
use mem::Mem;

use std::ptr;
use std::mem::{size_of, transmute};

pub struct CommandQueue {
    pub id: ll::CommandQueue,
}

pub struct ReadLock<'a, T: Sized + 'a> {
    #[allow(dead_code)]
    sl: &'a [T],
    evt: Event,
}

pub struct WriteLock<'a, T: Sized + 'a> {
    #[allow(dead_code)]
    sl: &'a mut [T],
    evt: Event,
}

impl<'a, T: Sized> Drop for ReadLock<'a, T> {
    fn drop(self: &mut Self) -> () {
        self.evt.wait().ok().expect(
            "error waiting to finish write for ReadLock");
    }
}

impl<'a, T: Sized> EventLike for ReadLock<'a, T> {
    fn status(self: &Self) -> Result<EventStatus, Error> {
        self.evt.status()
    }

    fn wait(self: &Self) -> Result<(), Error> {
        self.evt.wait()
    }
}

impl<'a, T: Sized> Drop for WriteLock<'a, T> {
    fn drop(self: &mut Self) -> () {
        self.evt.wait().ok().expect(
            "error waiting to finish write for ReadLock");
    }
}

impl<'a, T: Sized> EventLike for WriteLock<'a, T> {
    fn status(self: &Self) -> Result<EventStatus, Error> {
        self.evt.status()
    }

    fn wait(self: &Self) -> Result<(), Error> {
        self.evt.wait()
    }
}

unsafe impl Send for CommandQueue { }

impl CommandQueue {
    pub fn new(context: Context,
               device: Device) -> Result<CommandQueue, Error> {
        let props: c_ulong = if try!(device.out_of_order_supported()) {
            1
        } else {
            0
        };
        let id = unsafe {
            let mut err: i32 = 0;
            let id = ll::clCreateCommandQueue(context.id,
                                              device.id,
                                              props,
                                              &mut err);
            try!(Error::check(err));
            id
        };
        Ok(CommandQueue{ id: id })
    }

    fn get_info(self: &Self, param: u32) -> Result<*mut c_void, Error> {
        let mut tr: *mut c_void = ptr::null_mut();
        unsafe {
            try!(Error::check(ll::clGetCommandQueueInfo(self.id,
                                                   param,
                                                   size_of::<*mut c_void>() as size_t,
                                                   transmute(&mut tr),
                                                   ptr::null_mut())));
        }
        Ok(tr)
    }

    pub fn context(self: &Self) -> Result<Context, Error> {
        let val = try!(self.get_info(0x1090));
        let id = unsafe { transmute(val) };
        unsafe {
            try!(Error::check(ll::clRetainContext(id)));
        }
        Ok(Context{id: id})
    }

    pub fn device(self: &Self) -> Result<Device, Error> {
        let val = try!(self.get_info(0x1091));
        let id = unsafe { transmute(val) };
        Ok(Device{id: id})
    }

    fn alloc_ll_buffer(self: &Self,
                    flags: c_ulong,
                    size_bytes: usize) -> Result<ll::Mem, Error> {
        let context = try!(self.context());
        let id = unsafe {
            let mut err: i32 = 0;
            let id = ll::clCreateBuffer(context.id,
                                        flags,
                                        size_bytes as size_t,
                                        ptr::null_mut(),
                                        &mut err);
            try!(Error::check(err));
            id
        };
        Ok(id)
    }

    fn write_buffer_raw<T: Sized>(self: &Self,
                       buf: &mut Mem,
                       offset: usize,
                       size: usize,
                       data: &[T]) -> Result<Event, Error> {
        let mut event_id: ll::Event = ptr::null_mut();
        let host_ptr = &data[0] as *const T;

        unsafe {
            try!(Error::check(ll::clEnqueueWriteBuffer(self.id,
                                                       buf.id,
                                                       0, // don't block
                                                       offset as size_t,
                                                       size as size_t,
                                                       transmute(host_ptr),
                                                       0,
                                                       ptr::null(),
                                                       &mut event_id)));
        }

        let event = Event{id: event_id};
        Ok(event)
    }

    fn read_buffer_raw<T: Sized>(self: &Self,
                                 buf: &Mem,
                                 offset: usize,
                                 size: usize,
                                 data: &mut [T]) -> Result<Event, Error> {
        let mut event_id: ll::Event = ptr::null_mut();
        let host_ptr = &mut data[0] as *mut T;

        unsafe {
            try!(Error::check(ll::clEnqueueReadBuffer(self.id,
                                                      buf.id,
                                                      0, // don't block,
                                                      offset as size_t,
                                                      size as size_t,
                                                      transmute(host_ptr),
                                                      0,
                                                      ptr::null(),
                                                      &mut event_id)));
        }

        let event = Event{id: event_id};
        Ok(event)
    }

    pub fn create_buffer(self: &Self,
                               size_bytes: usize) -> Result<Mem, Error> {
        let id = try!(self.alloc_ll_buffer(0, size_bytes));
        let tr = Mem::new(id);
        Ok(tr)
    }

    pub fn create_buffer_from_slice<T: Sized>(self: &Self,
                                    slice: &[T]) -> Result<Mem, Error> {
        let num_bytes = slice.len() * size_of::<T>();
        let mut buf = try!(self.create_buffer(num_bytes));
        try!(self.write_buffer(&mut buf, slice));
        Ok(buf)
    }

    pub fn write_buffer_unscoped<T: Sized>(self: &Self,
                                  mem: &mut Mem,
                                  slice: &[T]) -> Result<Event, Error> {
        let size_bytes = size_of::<T>()*slice.len();
        let evt = try!(self.write_buffer_raw(mem,
                                             0,
                                             size_bytes,
                                             slice));
        Ok(evt)
    }

    pub fn read_buffer_unscoped<T: Sized>(self: &Self,
                                                 mem: &Mem,
                                                 slice: &mut [T]) -> Result<Event, Error> {
        let size_bytes = size_of::<T>()*slice.len();
        let evt = try!(self.read_buffer_raw(mem,
                                            0,
                                            size_bytes,
                                            slice));
        Ok(evt)
    }

    pub fn write_buffer<'a, 'b, T: Sized + 'a>(self: &Self,
                                  mem: &'b mut Mem,
                                  slice: &'a [T]) -> Result<ReadLock<'a, T>, Error> {
        let size_bytes = size_of::<T>()*slice.len();
        let evt = try!(self.write_buffer_raw(mem,
                                             0,
                                             size_bytes,
                                             slice));
        Ok(ReadLock{ evt: evt, sl: slice })
    }

    pub fn read_buffer<'a, T: Sized>(self: &Self,
                                 mem: &Mem,
                                 slice: &'a mut [T]) -> Result<WriteLock<'a, T>, Error> {
        let size_bytes = size_of::<T>()*slice.len();
        let evt = try!(self.read_buffer_raw(mem,
                                            0,
                                            size_bytes,
                                            slice));
        Ok(WriteLock{ evt: evt, sl: slice })
    }

    fn next_mul(self: &Self, global: usize, local: usize) -> usize {
        if global % local == 0 {
            global
        } else {
            ((global/local)+1)*local
        }
    }

    pub fn run(
                self: &Self,
                kernel: &mut Kernel,
                local_size: (usize, usize, usize),
                global_size: (usize, usize, usize)) -> Result<Event, Error> {
        let global_size_fixed = (self.next_mul(global_size.0, local_size.0),
                                self.next_mul(global_size.1, local_size.1),
                                self.next_mul(global_size.2, local_size.2));

        let local_vec = vec!(local_size.0 as size_t,
                             local_size.1 as size_t,
                             local_size.2 as size_t);
        let global_vec = vec!(global_size_fixed.0 as size_t,
                              global_size_fixed.1 as size_t,
                              global_size_fixed.2 as size_t);

        let mut event_id: ll::Event = ptr::null_mut();
        unsafe {
            let events_val = ptr::null_mut();
            try!(Error::check(ll::clEnqueueNDRangeKernel(self.id,
                                                         kernel.id,
                                                         3,
                                                         ptr::null(),
                                                         &global_vec[0],
                                                         &local_vec[0],
                                                         0,
                                                         events_val,
                                                         &mut event_id)));
        }
        let event = Event{id: event_id};

        Ok(event)
    }

    pub fn run_with_events(
                self: &Self,
                kernel: &mut Kernel,
                local_size: (usize, usize, usize),
                global_size: (usize, usize, usize),
                events_in: &[Event]) -> Result<Event, Error> {
        let events: Vec<ll::Event> = events_in.iter().map(|e| e.id).collect();
        let global_size_fixed = (self.next_mul(global_size.0, local_size.0),
                                self.next_mul(global_size.1, local_size.1),
                                self.next_mul(global_size.2, local_size.2));

        let local_vec = vec!(local_size.0 as size_t,
                             local_size.1 as size_t,
                             local_size.2 as size_t);
        let global_vec = vec!(global_size_fixed.0 as size_t,
                              global_size_fixed.1 as size_t,
                              global_size_fixed.2 as size_t);

        let mut event_id: ll::Event = ptr::null_mut();
        unsafe {
            let events_val = if events.len() > 0 {
                transmute(&events[0])
            } else {
                ptr::null_mut()
            };
            try!(Error::check(ll::clEnqueueNDRangeKernel(self.id,
                                                         kernel.id,
                                                         3,
                                                         ptr::null(),
                                                         &global_vec[0],
                                                         &local_vec[0],
                                                         events.len() as u32,
                                                         events_val,
                                                         &mut event_id)));
        }
        let event = Event{id: event_id};

        Ok(event)
    }
}

impl Drop for CommandQueue {
    fn drop(self: &mut Self) -> () {
        unsafe {
            Error::check(ll::clReleaseCommandQueue(self.id))
                .ok().expect("Error releasing command queue");
        }
    }
}

impl Clone for CommandQueue {
    fn clone(self: &Self) -> Self {
        unsafe {
            Error::check(ll::clRetainCommandQueue(self.id))
                .ok().expect("Error retaining command queue");
        }
        CommandQueue{ id: self.id }
    }
}

<|endoftext|>