<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 serde_derive;
#[macro_use]
extern crate structopt_derive;
extern crate structopt;
extern crate csv;
extern crate mimir;
extern crate itertools;
extern crate mimirsbrunn;
#[macro_use]
extern crate log;

use std::rc::Rc;
use mimir::rubber::Rubber;
use mimirsbrunn::utils::{format_label, get_zip_codes_from_admins};
use mimir::MimirObject;
use mimirsbrunn::admin_geofinder::AdminGeoFinder;
use std::collections::HashMap;
use structopt::StructOpt;

const GLOBAL_STOP_INDEX_NAME: &'static str = "munin_global_stops";
const MAX_LAT: f64 = 90f64;
const MIN_LAT: f64 = -90f64;

const MAX_LON: f64 = 180f64;
const MIN_LON: f64 = -180f64;

#[derive(Debug, StructOpt)]
struct Args {
    /// NTFS stops.txt file.
    #[structopt(short = "i", long = "input")]
    input: String,
    /// Name of the dataset.
    #[structopt(short = "d", long = "dataset", default_value = "fr")]
    dataset: String,
    /// Elasticsearch parameters.
    #[structopt(short = "c", long = "connection-string",
                default_value = "http://localhost:9200/munin")]
    connection_string: String,
    /// City level to calculate weight.
    #[structopt(short = "C", long = "city-level", default_value = "8")]
    city_level: u32,
}

#[derive(Deserialize, Debug)]
enum StopConvErr {
    InvisibleStop,
    StopPoint,
    InvalidStop(String),
}

#[derive(Debug, Deserialize)]
struct GtfsStop {
    stop_id: String,
    stop_lat: f64,
    stop_lon: f64,
    stop_name: String,
    location_type: Option<i32>,
    visible: Option<i32>,
    parent_station: Option<String>,
}

impl GtfsStop {
    fn incr_stop_point(&self, nb_stop_points: &mut HashMap<String, u32>) {
        match (self.location_type, &self.parent_station) {
            (Some(0), &Some(ref id)) |
            (None, &Some(ref id)) if !id.is_empty() => {
                *nb_stop_points
                    .entry(format!("stop_area:{}", id))
                    .or_insert(0) += 1
            }
            _ => (),
        }
    }
    // to be moved when TryInto is stablilized
    fn try_into(self) -> Result<mimir::Stop, StopConvErr> {
        if self.location_type != Some(1) {
            Err(StopConvErr::StopPoint)
        } else if self.visible == Some(0){
        	Err(StopConvErr::InvisibleStop)
        } else if self.stop_lat <= MIN_LAT || self.stop_lat >= MAX_LAT || self.stop_lon <= MIN_LON || self.stop_lon >= MAX_LON {
         	//Here we return an error message
		  	Err(StopConvErr::InvalidStop(format!("Invalid lon or lat for stop {}", self.stop_name).into()))
        } else {
        	Ok(mimir::Stop {
                id: format!("stop_area:{}", self.stop_id), // prefix to match navitia's id
                coord: mimir::Coord::new(self.stop_lat, self.stop_lon),
                label: self.stop_name.clone(),
                weight: 0.,
                zip_codes: vec![],
                administrative_regions: vec![],
                name: self.stop_name,
                coverages: vec![],
            })
        }
    }
}

fn attach_stop(stop: &mut mimir::Stop, admins: Vec<Rc<mimir::Admin>>, city_level: u32) {
    stop.administrative_regions = admins;
    stop.label = format_label(&stop.administrative_regions, city_level, &stop.name);
    stop.zip_codes = get_zip_codes_from_admins(&stop.administrative_regions);
}

/// Attach the stops to administrative regions
///
/// The admins are loaded from Elasticsearch and stored in a quadtree
/// We attach a stop with all the admins that have a boundary containing
/// the coordinate of the stop
fn attach_stops_to_admins<'a, It: Iterator<Item = &'a mut mimir::Stop>>(
    stops: It,
    rubber: &mut Rubber,
    city_level: u32,
) {
    let admins = rubber.get_all_admins().unwrap_or_else(|_| {
        info!("Administratives regions not found in elasticsearch db");
        vec![]
    });

    info!("{} administrative regions loaded from mimir", admins.len());

    let admins_geofinder = admins.into_iter().collect::<AdminGeoFinder>();

    let mut nb_unmatched = 0u32;
    let mut nb_matched = 0u32;
    for mut stop in stops {
        let admins = admins_geofinder.get(&stop.coord);

        if admins.is_empty() {
            nb_unmatched += 1;
        } else {
            nb_matched += 1;
        }

        attach_stop(&mut stop, admins, city_level);
    }

    info!(
        "there is {}/{} stops without any admin",
        nb_unmatched,
        nb_matched + nb_unmatched
    );
}

// Update weight value for each stop_area from HashMap and stop's coverage
fn finalize_stop_area<'a, It: Iterator<Item = &'a mut mimir::Stop>>(
    stops: It,
    nb_stop_points: &HashMap<String, u32>,
    dataset: &str,
) {
    let max = *nb_stop_points.values().max().unwrap_or(&1) as f64;
    for stop in stops {
        if let Some(weight) = nb_stop_points.get(&stop.id) {
            stop.weight = *weight as f64 / max;
        }
        stop.coverages.push(dataset.to_string());
    }
}

/// merge the stops from all the different indexes
/// for the moment the merge is very simple and uses only the ID
/// (and we take the data from the first stop inserted)
fn merge_stops<It: IntoIterator<Item = mimir::Stop>>(
    stops: It,
) -> Box<Iterator<Item = mimir::Stop>> {
    let mut stops_by_id = HashMap::<String, mimir::Stop>::new();
    for mut stop in stops.into_iter() {
        let cov = std::mem::replace(&mut stop.coverages, vec![]);
        stops_by_id
            .entry(stop.id.clone())
            .or_insert(stop)
            .coverages
            .extend(cov.into_iter());
    }
    Box::new(stops_by_id.into_iter().map(|(_, v)| v))
}

fn get_all_stops(rubber: &mut Rubber, index: String) -> Result<Vec<mimir::Stop>, String> {
    rubber.get_all_objects_from_index(&index).map_err(
        |e| e.to_string(),
    )
}

fn update_global_stop_index<'a, It: Iterator<Item = &'a mimir::Stop>>(
    rubber: &mut Rubber,
    stops: It,
    dataset: &str,
) -> Result<String, String> {
    let dataset_index =
        mimir::rubber::get_main_type_and_dataset_index(mimir::Stop::doc_type(), dataset);
    let stops_indexes = rubber
        .get_all_aliased_index(&mimir::rubber::get_main_type_index(mimir::Stop::doc_type()))?
        .into_iter()
        .filter(|&(_, ref aliases)| !aliases.contains(&dataset_index))
        .map(|(index, _)| index);

    let all_other_es_stops: Vec<_> = stops_indexes
        .map(|index| get_all_stops(rubber, index).unwrap())
        .flat_map(|stops| stops.into_iter())
        .collect();

    let all_es_stops = all_other_es_stops.into_iter().chain(
        stops.into_iter().cloned(),
    );

    let all_merged_stops = merge_stops(all_es_stops);
    let es_index_name = mimir::rubber::get_date_index_name(GLOBAL_STOP_INDEX_NAME);

    rubber.create_index(&es_index_name)?;

    let nb_stops_added = rubber
        .bulk_index(&es_index_name, all_merged_stops)
        .map_err(|e| e.to_string())?;
    info!("{} stops added in the global index", nb_stops_added);
    // create global index
    // fill structure for each stop indexes
    Ok(es_index_name)
}

// publish the global stop index
// alias the new index to the global stop alias, and remove the old index
fn publish_global_index(rubber: &mut Rubber, new_global_index: &str) -> Result<(), String> {
    let last_global_indexes: Vec<_> = rubber
        .get_all_aliased_index(GLOBAL_STOP_INDEX_NAME)?
        .into_iter()
        .map(|(k, _)| k)
        .filter(|k| k != new_global_index)
        .collect();
    rubber.alias(
        GLOBAL_STOP_INDEX_NAME,
        &vec![new_global_index.to_string()],
        &last_global_indexes,
    )?;

    for index in last_global_indexes {
        rubber.delete_index(&index)?;
    }
    Ok(())
}

fn main() {
    mimir::logger_init().unwrap();
    info!("Launching stops2mimir...");

    let args = Args::from_args();

    info!("creation of indexes");
    let mut rubber = Rubber::new(&args.connection_string);
    let mut rdr = csv::Reader::from_path(args.input).unwrap();

    let mut nb_stop_points = HashMap::new();

    let mut stops: Vec<mimir::Stop> = rdr.deserialize()
        .filter_map(|rc| {
            rc.map_err(|e| warn!("skip csv line because: {}", e)).ok()
        })
        .filter_map(|stop: GtfsStop| {
            stop.incr_stop_point(&mut nb_stop_points);
            match stop.try_into() {
                Ok(s) => Some(s),
                Err(StopConvErr::InvisibleStop) | Err(StopConvErr::StopPoint) => None,
                Err(StopConvErr::InvalidStop(msg)) => {
                    warn!("skip csv line because: {}", msg);
                    None
                }
            }
        })
        .collect();

    attach_stops_to_admins(stops.iter_mut(), &mut rubber, args.city_level);

    finalize_stop_area(stops.iter_mut(), &nb_stop_points, &args.dataset);

    let global_index = update_global_stop_index(&mut rubber, stops.iter(), &args.dataset).unwrap();

    info!("Importing {} stops into Mimir", stops.len());
    let nb_stops = rubber.index(&args.dataset, stops.iter()).unwrap();
    info!("Nb of indexed stops: {}", nb_stops);

    publish_global_index(&mut rubber, &global_index).unwrap();
}

#[test]
fn test_load_stops() {
    use itertools::Itertools;
    let mut rdr = csv::Reader::from_path("./tests/fixtures/stops.txt".to_string()).unwrap();

    let mut nb_stop_points = HashMap::new();
    let stops: Vec<mimir::Stop> = rdr.deserialize()
        .filter_map(Result::ok)
        .filter_map(|stop: GtfsStop| {
            stop.incr_stop_point(&mut nb_stop_points);
            stop.try_into()
        })
        .collect();
    let ids: Vec<_> = stops.iter().map(|s| s.id.clone()).sorted();
    assert_eq!(
        ids,
        vec![
            "stop_area:SA:known_by_all_dataset",
            "stop_area:SA:main_station",
            "stop_area:SA:second_station",
            "stop_area:SA:station_no_city",
            "stop_area:SA:weight_1_station",
            "stop_area:SA:weight_3_station",
        ]
    );
    let weights: Vec<_> = ids.iter().map(|id| nb_stop_points.get(id)).collect();
    assert_eq!(
        weights,
        vec![None, Some(&1), Some(&1), None, Some(&1), Some(&3)]
    );
}
<commit_msg>stops2mimir: modification applied in the test<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 serde_derive;
#[macro_use]
extern crate structopt_derive;
extern crate structopt;
extern crate csv;
extern crate mimir;
extern crate itertools;
extern crate mimirsbrunn;
#[macro_use]
extern crate log;

use std::rc::Rc;
use mimir::rubber::Rubber;
use mimirsbrunn::utils::{format_label, get_zip_codes_from_admins};
use mimir::MimirObject;
use mimirsbrunn::admin_geofinder::AdminGeoFinder;
use std::collections::HashMap;
use structopt::StructOpt;

const GLOBAL_STOP_INDEX_NAME: &'static str = "munin_global_stops";
const MAX_LAT: f64 = 90f64;
const MIN_LAT: f64 = -90f64;

const MAX_LON: f64 = 180f64;
const MIN_LON: f64 = -180f64;

#[derive(Debug, StructOpt)]
struct Args {
    /// NTFS stops.txt file.
    #[structopt(short = "i", long = "input")]
    input: String,
    /// Name of the dataset.
    #[structopt(short = "d", long = "dataset", default_value = "fr")]
    dataset: String,
    /// Elasticsearch parameters.
    #[structopt(short = "c", long = "connection-string",
                default_value = "http://localhost:9200/munin")]
    connection_string: String,
    /// City level to calculate weight.
    #[structopt(short = "C", long = "city-level", default_value = "8")]
    city_level: u32,
}

#[derive(Deserialize, Debug)]
enum StopConvErr {
    InvisibleStop,
    StopPoint,
    InvalidStop(String),
}

#[derive(Debug, Deserialize)]
struct GtfsStop {
    stop_id: String,
    stop_lat: f64,
    stop_lon: f64,
    stop_name: String,
    location_type: Option<i32>,
    visible: Option<i32>,
    parent_station: Option<String>,
}

impl GtfsStop {
    fn incr_stop_point(&self, nb_stop_points: &mut HashMap<String, u32>) {
        match (self.location_type, &self.parent_station) {
            (Some(0), &Some(ref id)) |
            (None, &Some(ref id)) if !id.is_empty() => {
                *nb_stop_points
                    .entry(format!("stop_area:{}", id))
                    .or_insert(0) += 1
            }
            _ => (),
        }
    }
    // to be moved when TryInto is stablilized
    fn try_into(self) -> Result<mimir::Stop, StopConvErr> {
        if self.location_type != Some(1) {
            Err(StopConvErr::StopPoint)
        } else if self.visible == Some(0){
        	Err(StopConvErr::InvisibleStop)
        } else if self.stop_lat <= MIN_LAT || self.stop_lat >= MAX_LAT || self.stop_lon <= MIN_LON || self.stop_lon >= MAX_LON {
         	//Here we return an error message
		  	Err(StopConvErr::InvalidStop(format!("Invalid lon or lat for stop {}", self.stop_name).into()))
        } else {
        	Ok(mimir::Stop {
                id: format!("stop_area:{}", self.stop_id), // prefix to match navitia's id
                coord: mimir::Coord::new(self.stop_lat, self.stop_lon),
                label: self.stop_name.clone(),
                weight: 0.,
                zip_codes: vec![],
                administrative_regions: vec![],
                name: self.stop_name,
                coverages: vec![],
            })
        }
    }
}

fn attach_stop(stop: &mut mimir::Stop, admins: Vec<Rc<mimir::Admin>>, city_level: u32) {
    stop.administrative_regions = admins;
    stop.label = format_label(&stop.administrative_regions, city_level, &stop.name);
    stop.zip_codes = get_zip_codes_from_admins(&stop.administrative_regions);
}

/// Attach the stops to administrative regions
///
/// The admins are loaded from Elasticsearch and stored in a quadtree
/// We attach a stop with all the admins that have a boundary containing
/// the coordinate of the stop
fn attach_stops_to_admins<'a, It: Iterator<Item = &'a mut mimir::Stop>>(
    stops: It,
    rubber: &mut Rubber,
    city_level: u32,
) {
    let admins = rubber.get_all_admins().unwrap_or_else(|_| {
        info!("Administratives regions not found in elasticsearch db");
        vec![]
    });

    info!("{} administrative regions loaded from mimir", admins.len());

    let admins_geofinder = admins.into_iter().collect::<AdminGeoFinder>();

    let mut nb_unmatched = 0u32;
    let mut nb_matched = 0u32;
    for mut stop in stops {
        let admins = admins_geofinder.get(&stop.coord);

        if admins.is_empty() {
            nb_unmatched += 1;
        } else {
            nb_matched += 1;
        }

        attach_stop(&mut stop, admins, city_level);
    }

    info!(
        "there is {}/{} stops without any admin",
        nb_unmatched,
        nb_matched + nb_unmatched
    );
}

// Update weight value for each stop_area from HashMap and stop's coverage
fn finalize_stop_area<'a, It: Iterator<Item = &'a mut mimir::Stop>>(
    stops: It,
    nb_stop_points: &HashMap<String, u32>,
    dataset: &str,
) {
    let max = *nb_stop_points.values().max().unwrap_or(&1) as f64;
    for stop in stops {
        if let Some(weight) = nb_stop_points.get(&stop.id) {
            stop.weight = *weight as f64 / max;
        }
        stop.coverages.push(dataset.to_string());
    }
}

/// merge the stops from all the different indexes
/// for the moment the merge is very simple and uses only the ID
/// (and we take the data from the first stop inserted)
fn merge_stops<It: IntoIterator<Item = mimir::Stop>>(
    stops: It,
) -> Box<Iterator<Item = mimir::Stop>> {
    let mut stops_by_id = HashMap::<String, mimir::Stop>::new();
    for mut stop in stops.into_iter() {
        let cov = std::mem::replace(&mut stop.coverages, vec![]);
        stops_by_id
            .entry(stop.id.clone())
            .or_insert(stop)
            .coverages
            .extend(cov.into_iter());
    }
    Box::new(stops_by_id.into_iter().map(|(_, v)| v))
}

fn get_all_stops(rubber: &mut Rubber, index: String) -> Result<Vec<mimir::Stop>, String> {
    rubber.get_all_objects_from_index(&index).map_err(
        |e| e.to_string(),
    )
}

fn update_global_stop_index<'a, It: Iterator<Item = &'a mimir::Stop>>(
    rubber: &mut Rubber,
    stops: It,
    dataset: &str,
) -> Result<String, String> {
    let dataset_index =
        mimir::rubber::get_main_type_and_dataset_index(mimir::Stop::doc_type(), dataset);
    let stops_indexes = rubber
        .get_all_aliased_index(&mimir::rubber::get_main_type_index(mimir::Stop::doc_type()))?
        .into_iter()
        .filter(|&(_, ref aliases)| !aliases.contains(&dataset_index))
        .map(|(index, _)| index);

    let all_other_es_stops: Vec<_> = stops_indexes
        .map(|index| get_all_stops(rubber, index).unwrap())
        .flat_map(|stops| stops.into_iter())
        .collect();

    let all_es_stops = all_other_es_stops.into_iter().chain(
        stops.into_iter().cloned(),
    );

    let all_merged_stops = merge_stops(all_es_stops);
    let es_index_name = mimir::rubber::get_date_index_name(GLOBAL_STOP_INDEX_NAME);

    rubber.create_index(&es_index_name)?;

    let nb_stops_added = rubber
        .bulk_index(&es_index_name, all_merged_stops)
        .map_err(|e| e.to_string())?;
    info!("{} stops added in the global index", nb_stops_added);
    // create global index
    // fill structure for each stop indexes
    Ok(es_index_name)
}

// publish the global stop index
// alias the new index to the global stop alias, and remove the old index
fn publish_global_index(rubber: &mut Rubber, new_global_index: &str) -> Result<(), String> {
    let last_global_indexes: Vec<_> = rubber
        .get_all_aliased_index(GLOBAL_STOP_INDEX_NAME)?
        .into_iter()
        .map(|(k, _)| k)
        .filter(|k| k != new_global_index)
        .collect();
    rubber.alias(
        GLOBAL_STOP_INDEX_NAME,
        &vec![new_global_index.to_string()],
        &last_global_indexes,
    )?;

    for index in last_global_indexes {
        rubber.delete_index(&index)?;
    }
    Ok(())
}

fn main() {
    mimir::logger_init().unwrap();
    info!("Launching stops2mimir...");

    let args = Args::from_args();

    info!("creation of indexes");
    let mut rubber = Rubber::new(&args.connection_string);
    let mut rdr = csv::Reader::from_path(args.input).unwrap();

    let mut nb_stop_points = HashMap::new();

    let mut stops: Vec<mimir::Stop> = rdr.deserialize()
        .filter_map(|rc| {
            rc.map_err(|e| warn!("skip csv line because: {}", e)).ok()
        })
        .filter_map(|stop: GtfsStop| {
            stop.incr_stop_point(&mut nb_stop_points);
            match stop.try_into() {
                Ok(s) => Some(s),
                Err(StopConvErr::InvisibleStop) | Err(StopConvErr::StopPoint) => None,
                Err(StopConvErr::InvalidStop(msg)) => {
                    warn!("skip csv line because: {}", msg);
                    None
                }
            }
        })
        .collect();

    attach_stops_to_admins(stops.iter_mut(), &mut rubber, args.city_level);

    finalize_stop_area(stops.iter_mut(), &nb_stop_points, &args.dataset);

    let global_index = update_global_stop_index(&mut rubber, stops.iter(), &args.dataset).unwrap();

    info!("Importing {} stops into Mimir", stops.len());
    let nb_stops = rubber.index(&args.dataset, stops.iter()).unwrap();
    info!("Nb of indexed stops: {}", nb_stops);

    publish_global_index(&mut rubber, &global_index).unwrap();
}

#[test]
fn test_load_stops() {
    use itertools::Itertools;
    let mut rdr = csv::Reader::from_path("./tests/fixtures/stops.txt".to_string()).unwrap();

    let mut nb_stop_points = HashMap::new();
    let stops: Vec<mimir::Stop> = rdr.deserialize()
        .filter_map(Result::ok)
        .filter_map(|stop: GtfsStop| {
            stop.incr_stop_point(&mut nb_stop_points);
            match stop.try_into() {
                Ok(s) => Some(s),
                Err(StopConvErr::InvisibleStop) => None,
                Err(StopConvErr::StopPoint) => None,
                Err(StopConvErr::InvalidStop(msg)) => {
                    warn!("skip csv line because: {}", msg);
                    None
                }
            }
        })
        .collect();
    let ids: Vec<_> = stops.iter().map(|s| s.id.clone()).sorted();
    assert_eq!(
        ids,
        vec![
            "stop_area:SA:known_by_all_dataset",
            "stop_area:SA:main_station",
            "stop_area:SA:second_station",
            "stop_area:SA:station_no_city",
            "stop_area:SA:weight_1_station",
            "stop_area:SA:weight_3_station",
        ]
    );
    let weights: Vec<_> = ids.iter().map(|id| nb_stop_points.get(id)).collect();
    assert_eq!(
        weights,
        vec![None, Some(&1), Some(&1), None, Some(&1), Some(&3)]
    );
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use std::collections::HashMap;
use std::mem;

use bindgen::bindings::Bindings;
use bindgen::config::{Config, Language};
use bindgen::declarationtyperesolver::DeclarationTypeResolver;
use bindgen::dependencies::Dependencies;
use bindgen::error::Error;
use bindgen::ir::{Constant, Enum, Function, Item, ItemContainer, ItemMap};
use bindgen::ir::{OpaqueItem, Path, Static, Struct, Typedef, Union};
use bindgen::monomorph::Monomorphs;
use bindgen::ItemType;

#[derive(Debug, Clone)]
pub struct Library {
    config: Config,
    constants: ItemMap<Constant>,
    globals: ItemMap<Static>,
    enums: ItemMap<Enum>,
    structs: ItemMap<Struct>,
    unions: ItemMap<Union>,
    opaque_items: ItemMap<OpaqueItem>,
    typedefs: ItemMap<Typedef>,
    functions: Vec<Function>,
}

impl Library {
    pub fn new(
        config: Config,
        constants: ItemMap<Constant>,
        globals: ItemMap<Static>,
        enums: ItemMap<Enum>,
        structs: ItemMap<Struct>,
        unions: ItemMap<Union>,
        opaque_items: ItemMap<OpaqueItem>,
        typedefs: ItemMap<Typedef>,
        functions: Vec<Function>,
    ) -> Library {
        Library {
            config: config,
            constants: constants,
            globals: globals,
            enums: enums,
            structs: structs,
            unions: unions,
            opaque_items: opaque_items,
            typedefs: typedefs,
            functions: functions,
        }
    }

    pub fn generate(mut self) -> Result<Bindings, Error> {
        self.remove_excluded();
        self.functions.sort_by(|x, y| x.name.cmp(&y.name));
        self.transfer_annotations();
        self.rename_items();
        self.simplify_option_to_ptr();

        if self.config.language == Language::C {
            self.instantiate_monomorphs();

            self.resolve_declaration_types();
        }

        let mut dependencies = Dependencies::new();

        for function in &self.functions {
            function.add_dependencies(&self, &mut dependencies);
        }
        self.globals.for_all_items(|global| {
            global.add_dependencies(&self, &mut dependencies);
        });
        for name in &self.config.export.include {
            if let Some(items) = self.get_items(name) {
                if !dependencies.items.contains(name) {
                    dependencies.items.insert(name.clone());

                    for item in &items {
                        item.deref().add_dependencies(&self, &mut dependencies);
                    }
                    for item in items {
                        dependencies.order.push(item);
                    }
                }
            }
        }

        dependencies.sort();

        let items = dependencies.order;
        let constants =
            if self.config.export.should_generate(ItemType::Constants) {
                self.constants.to_vec()
            } else {
                vec![]
            };

        let globals =
            if self.config.export.should_generate(ItemType::Globals) {
                self.globals.to_vec()
            } else {
                vec![]
            };
        let functions =
            if self.config.export.should_generate(ItemType::Functions) {
                mem::replace(&mut self.functions, vec![])
            } else {
                vec![]
            };

        Ok(Bindings::new(
            self.config,
            constants,
            globals,
            items,
            functions,
        ))
    }

    pub fn get_items(&self, p: &Path) -> Option<Vec<ItemContainer>> {
        macro_rules! find {
            ($field:ident, $kind:ident) => {
                if self.config.export.should_generate(ItemType::$kind) {
                    if let Some(x) = self.$field.get_items(p) {
                        return Some(x)
                    }
                }
            }
        }

        find!(enums, Enums);
        find!(structs, Structs);
        find!(unions, Unions);
        find!(opaque_items, OpaqueItems);
        find!(typedefs, Typedefs);

        None
    }

    fn remove_excluded(&mut self) {
        let config = &self.config;
        self.functions
            .retain(|x| !config.export.exclude.contains(&x.name));
        self.enums
            .filter(|x| config.export.exclude.contains(&x.name));
        self.structs
            .filter(|x| config.export.exclude.contains(&x.name));
        self.unions
            .filter(|x| config.export.exclude.contains(&x.name));
        self.opaque_items
            .filter(|x| config.export.exclude.contains(&x.name));
        self.typedefs
            .filter(|x| config.export.exclude.contains(&x.name));
        self.globals
            .filter(|x| config.export.exclude.contains(&x.name));
        self.constants
            .filter(|x| config.export.exclude.contains(&x.name));
    }

    fn transfer_annotations(&mut self) {
        let mut annotations = HashMap::new();

        self.typedefs.for_all_items_mut(|x| {
            x.transfer_annotations(&mut annotations);
        });

        for (alias_path, annotations) in annotations {
            // TODO
            let mut transferred = false;

            self.enums.for_items_mut(&alias_path, |x| {
                if x.annotations().is_empty() {
                    *x.annotations_mut() = annotations.clone();
                    transferred = true;
                } else {
                    warn!(
                        "Can't transfer annotations from typedef to alias ({}) \
                         that already has annotations.",
                        alias_path
                    );
                }
            });
            if transferred {
                continue;
            }
            self.structs.for_items_mut(&alias_path, |x| {
                if x.annotations().is_empty() {
                    *x.annotations_mut() = annotations.clone();
                    transferred = true;
                } else {
                    warn!(
                        "Can't transfer annotations from typedef to alias ({}) \
                         that already has annotations.",
                        alias_path
                    );
                }
            });
            if transferred {
                continue;
            }
            self.unions.for_items_mut(&alias_path, |x| {
                if x.annotations().is_empty() {
                    *x.annotations_mut() = annotations.clone();
                    transferred = true;
                } else {
                    warn!(
                        "Can't transfer annotations from typedef to alias ({}) \
                         that already has annotations.",
                        alias_path
                    );
                }
            });
            if transferred {
                continue;
            }
            self.opaque_items.for_items_mut(&alias_path, |x| {
                if x.annotations().is_empty() {
                    *x.annotations_mut() = annotations.clone();
                    transferred = true;
                } else {
                    warn!(
                        "Can't transfer annotations from typedef to alias ({}) \
                         that already has annotations.",
                        alias_path
                    );
                }
            });
            if transferred {
                continue;
            }
            self.typedefs.for_items_mut(&alias_path, |x| {
                if x.annotations().is_empty() {
                    *x.annotations_mut() = annotations.clone();
                    transferred = true;
                } else {
                    warn!(
                        "Can't transfer annotations from typedef to alias ({}) \
                         that already has annotations.",
                        alias_path
                    );
                }
            });
            if transferred {
                continue;
            }
        }
    }

    fn rename_items(&mut self) {
        let config = &self.config;

        self.globals
            .for_all_items_mut(|x| x.rename_for_config(config));
        self.globals.rebuild();

        self.constants
            .for_all_items_mut(|x| x.rename_for_config(config));
        self.constants.rebuild();

        self.structs
            .for_all_items_mut(|x| x.rename_for_config(config));
        self.structs.rebuild();

        self.unions
            .for_all_items_mut(|x| x.rename_for_config(config));
        self.unions.rebuild();

        self.enums
            .for_all_items_mut(|x| x.rename_for_config(config));
        self.enums.rebuild();

        self.opaque_items
            .for_all_items_mut(|x| x.rename_for_config(config));
        self.opaque_items.rebuild();

        self.typedefs
            .for_all_items_mut(|x| x.rename_for_config(config));
        self.typedefs.rebuild();

        for item in &mut self.functions {
            item.rename_for_config(&self.config);
        }
    }

    fn resolve_declaration_types(&mut self) {
        if self.config.style.generate_typedef() {
            return;
        }

        let mut resolver = DeclarationTypeResolver::new();

        self.structs.for_all_items(|x| {
            x.collect_declaration_types(&mut resolver);
        });

        self.opaque_items.for_all_items(|x| {
            x.collect_declaration_types(&mut resolver);
        });

        self.enums.for_all_items(|x| {
            x.collect_declaration_types(&mut resolver);
        });

        self.unions.for_all_items(|x| {
            x.collect_declaration_types(&mut resolver);
        });

        self.enums
            .for_all_items_mut(|x| x.resolve_declaration_types(&resolver));

        self.structs
            .for_all_items_mut(|x| x.resolve_declaration_types(&resolver));

        self.unions
            .for_all_items_mut(|x| x.resolve_declaration_types(&resolver));

        self.typedefs
            .for_all_items_mut(|x| x.resolve_declaration_types(&resolver));

        self.globals
            .for_all_items_mut(|x| x.resolve_declaration_types(&resolver));

        for item in &mut self.functions {
            item.resolve_declaration_types(&resolver);
        }
    }

    fn simplify_option_to_ptr(&mut self) {
        self.structs.for_all_items_mut(|x| {
            x.simplify_option_to_ptr();
        });
        self.unions.for_all_items_mut(|x| {
            x.simplify_option_to_ptr();
        });
        self.globals.for_all_items_mut(|x| {
            x.simplify_option_to_ptr();
        });
        self.typedefs.for_all_items_mut(|x| {
            x.simplify_option_to_ptr();
        });
        for x in &mut self.functions {
            x.simplify_option_to_ptr();
        }
    }

    fn instantiate_monomorphs(&mut self) {
        // Collect a list of monomorphs
        let mut monomorphs = Monomorphs::default();

        self.structs.for_all_items(|x| {
            x.add_monomorphs(self, &mut monomorphs);
        });
        self.unions.for_all_items(|x| {
            x.add_monomorphs(self, &mut monomorphs);
        });
        self.typedefs.for_all_items(|x| {
            x.add_monomorphs(self, &mut monomorphs);
        });
        for x in &self.functions {
            x.add_monomorphs(self, &mut monomorphs);
        }

        // Insert the monomorphs into self
        for monomorph in monomorphs.drain_structs() {
            self.structs.try_insert(monomorph);
        }
        for monomorph in monomorphs.drain_unions() {
            self.unions.try_insert(monomorph);
        }
        for monomorph in monomorphs.drain_opaques() {
            self.opaque_items.try_insert(monomorph);
        }
        for monomorph in monomorphs.drain_typedefs() {
            self.typedefs.try_insert(monomorph);
        }

        // Remove structs and opaque items that are generic
        self.opaque_items.filter(|x| x.generic_params.len() > 0);
        self.structs.filter(|x| x.generic_params.len() > 0);
        self.unions.filter(|x| x.generic_params.len() > 0);
        self.typedefs.filter(|x| x.generic_params.len() > 0);

        // Mangle the paths that remain
        self.unions
            .for_all_items_mut(|x| x.mangle_paths(&monomorphs));
        self.structs
            .for_all_items_mut(|x| x.mangle_paths(&monomorphs));
        self.typedefs
            .for_all_items_mut(|x| x.mangle_paths(&monomorphs));
        for x in &mut self.functions {
            x.mangle_paths(&monomorphs);
        }
    }
}
<commit_msg>Rustfmt.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use std::collections::HashMap;
use std::mem;

use bindgen::bindings::Bindings;
use bindgen::config::{Config, Language};
use bindgen::declarationtyperesolver::DeclarationTypeResolver;
use bindgen::dependencies::Dependencies;
use bindgen::error::Error;
use bindgen::ir::{Constant, Enum, Function, Item, ItemContainer, ItemMap};
use bindgen::ir::{OpaqueItem, Path, Static, Struct, Typedef, Union};
use bindgen::monomorph::Monomorphs;
use bindgen::ItemType;

#[derive(Debug, Clone)]
pub struct Library {
    config: Config,
    constants: ItemMap<Constant>,
    globals: ItemMap<Static>,
    enums: ItemMap<Enum>,
    structs: ItemMap<Struct>,
    unions: ItemMap<Union>,
    opaque_items: ItemMap<OpaqueItem>,
    typedefs: ItemMap<Typedef>,
    functions: Vec<Function>,
}

impl Library {
    pub fn new(
        config: Config,
        constants: ItemMap<Constant>,
        globals: ItemMap<Static>,
        enums: ItemMap<Enum>,
        structs: ItemMap<Struct>,
        unions: ItemMap<Union>,
        opaque_items: ItemMap<OpaqueItem>,
        typedefs: ItemMap<Typedef>,
        functions: Vec<Function>,
    ) -> Library {
        Library {
            config: config,
            constants: constants,
            globals: globals,
            enums: enums,
            structs: structs,
            unions: unions,
            opaque_items: opaque_items,
            typedefs: typedefs,
            functions: functions,
        }
    }

    pub fn generate(mut self) -> Result<Bindings, Error> {
        self.remove_excluded();
        self.functions.sort_by(|x, y| x.name.cmp(&y.name));
        self.transfer_annotations();
        self.rename_items();
        self.simplify_option_to_ptr();

        if self.config.language == Language::C {
            self.instantiate_monomorphs();

            self.resolve_declaration_types();
        }

        let mut dependencies = Dependencies::new();

        for function in &self.functions {
            function.add_dependencies(&self, &mut dependencies);
        }
        self.globals.for_all_items(|global| {
            global.add_dependencies(&self, &mut dependencies);
        });
        for name in &self.config.export.include {
            if let Some(items) = self.get_items(name) {
                if !dependencies.items.contains(name) {
                    dependencies.items.insert(name.clone());

                    for item in &items {
                        item.deref().add_dependencies(&self, &mut dependencies);
                    }
                    for item in items {
                        dependencies.order.push(item);
                    }
                }
            }
        }

        dependencies.sort();

        let items = dependencies.order;
        let constants = if self.config.export.should_generate(ItemType::Constants) {
            self.constants.to_vec()
        } else {
            vec![]
        };

        let globals = if self.config.export.should_generate(ItemType::Globals) {
            self.globals.to_vec()
        } else {
            vec![]
        };
        let functions = if self.config.export.should_generate(ItemType::Functions) {
            mem::replace(&mut self.functions, vec![])
        } else {
            vec![]
        };

        Ok(Bindings::new(
            self.config,
            constants,
            globals,
            items,
            functions,
        ))
    }

    pub fn get_items(&self, p: &Path) -> Option<Vec<ItemContainer>> {
        macro_rules! find {
            ($field:ident, $kind:ident) => {
                if self.config.export.should_generate(ItemType::$kind) {
                    if let Some(x) = self.$field.get_items(p) {
                        return Some(x);
                    }
                }
            };
        }

        find!(enums, Enums);
        find!(structs, Structs);
        find!(unions, Unions);
        find!(opaque_items, OpaqueItems);
        find!(typedefs, Typedefs);

        None
    }

    fn remove_excluded(&mut self) {
        let config = &self.config;
        self.functions
            .retain(|x| !config.export.exclude.contains(&x.name));
        self.enums
            .filter(|x| config.export.exclude.contains(&x.name));
        self.structs
            .filter(|x| config.export.exclude.contains(&x.name));
        self.unions
            .filter(|x| config.export.exclude.contains(&x.name));
        self.opaque_items
            .filter(|x| config.export.exclude.contains(&x.name));
        self.typedefs
            .filter(|x| config.export.exclude.contains(&x.name));
        self.globals
            .filter(|x| config.export.exclude.contains(&x.name));
        self.constants
            .filter(|x| config.export.exclude.contains(&x.name));
    }

    fn transfer_annotations(&mut self) {
        let mut annotations = HashMap::new();

        self.typedefs.for_all_items_mut(|x| {
            x.transfer_annotations(&mut annotations);
        });

        for (alias_path, annotations) in annotations {
            // TODO
            let mut transferred = false;

            self.enums.for_items_mut(&alias_path, |x| {
                if x.annotations().is_empty() {
                    *x.annotations_mut() = annotations.clone();
                    transferred = true;
                } else {
                    warn!(
                        "Can't transfer annotations from typedef to alias ({}) \
                         that already has annotations.",
                        alias_path
                    );
                }
            });
            if transferred {
                continue;
            }
            self.structs.for_items_mut(&alias_path, |x| {
                if x.annotations().is_empty() {
                    *x.annotations_mut() = annotations.clone();
                    transferred = true;
                } else {
                    warn!(
                        "Can't transfer annotations from typedef to alias ({}) \
                         that already has annotations.",
                        alias_path
                    );
                }
            });
            if transferred {
                continue;
            }
            self.unions.for_items_mut(&alias_path, |x| {
                if x.annotations().is_empty() {
                    *x.annotations_mut() = annotations.clone();
                    transferred = true;
                } else {
                    warn!(
                        "Can't transfer annotations from typedef to alias ({}) \
                         that already has annotations.",
                        alias_path
                    );
                }
            });
            if transferred {
                continue;
            }
            self.opaque_items.for_items_mut(&alias_path, |x| {
                if x.annotations().is_empty() {
                    *x.annotations_mut() = annotations.clone();
                    transferred = true;
                } else {
                    warn!(
                        "Can't transfer annotations from typedef to alias ({}) \
                         that already has annotations.",
                        alias_path
                    );
                }
            });
            if transferred {
                continue;
            }
            self.typedefs.for_items_mut(&alias_path, |x| {
                if x.annotations().is_empty() {
                    *x.annotations_mut() = annotations.clone();
                    transferred = true;
                } else {
                    warn!(
                        "Can't transfer annotations from typedef to alias ({}) \
                         that already has annotations.",
                        alias_path
                    );
                }
            });
            if transferred {
                continue;
            }
        }
    }

    fn rename_items(&mut self) {
        let config = &self.config;

        self.globals
            .for_all_items_mut(|x| x.rename_for_config(config));
        self.globals.rebuild();

        self.constants
            .for_all_items_mut(|x| x.rename_for_config(config));
        self.constants.rebuild();

        self.structs
            .for_all_items_mut(|x| x.rename_for_config(config));
        self.structs.rebuild();

        self.unions
            .for_all_items_mut(|x| x.rename_for_config(config));
        self.unions.rebuild();

        self.enums
            .for_all_items_mut(|x| x.rename_for_config(config));
        self.enums.rebuild();

        self.opaque_items
            .for_all_items_mut(|x| x.rename_for_config(config));
        self.opaque_items.rebuild();

        self.typedefs
            .for_all_items_mut(|x| x.rename_for_config(config));
        self.typedefs.rebuild();

        for item in &mut self.functions {
            item.rename_for_config(&self.config);
        }
    }

    fn resolve_declaration_types(&mut self) {
        if self.config.style.generate_typedef() {
            return;
        }

        let mut resolver = DeclarationTypeResolver::new();

        self.structs.for_all_items(|x| {
            x.collect_declaration_types(&mut resolver);
        });

        self.opaque_items.for_all_items(|x| {
            x.collect_declaration_types(&mut resolver);
        });

        self.enums.for_all_items(|x| {
            x.collect_declaration_types(&mut resolver);
        });

        self.unions.for_all_items(|x| {
            x.collect_declaration_types(&mut resolver);
        });

        self.enums
            .for_all_items_mut(|x| x.resolve_declaration_types(&resolver));

        self.structs
            .for_all_items_mut(|x| x.resolve_declaration_types(&resolver));

        self.unions
            .for_all_items_mut(|x| x.resolve_declaration_types(&resolver));

        self.typedefs
            .for_all_items_mut(|x| x.resolve_declaration_types(&resolver));

        self.globals
            .for_all_items_mut(|x| x.resolve_declaration_types(&resolver));

        for item in &mut self.functions {
            item.resolve_declaration_types(&resolver);
        }
    }

    fn simplify_option_to_ptr(&mut self) {
        self.structs.for_all_items_mut(|x| {
            x.simplify_option_to_ptr();
        });
        self.unions.for_all_items_mut(|x| {
            x.simplify_option_to_ptr();
        });
        self.globals.for_all_items_mut(|x| {
            x.simplify_option_to_ptr();
        });
        self.typedefs.for_all_items_mut(|x| {
            x.simplify_option_to_ptr();
        });
        for x in &mut self.functions {
            x.simplify_option_to_ptr();
        }
    }

    fn instantiate_monomorphs(&mut self) {
        // Collect a list of monomorphs
        let mut monomorphs = Monomorphs::default();

        self.structs.for_all_items(|x| {
            x.add_monomorphs(self, &mut monomorphs);
        });
        self.unions.for_all_items(|x| {
            x.add_monomorphs(self, &mut monomorphs);
        });
        self.typedefs.for_all_items(|x| {
            x.add_monomorphs(self, &mut monomorphs);
        });
        for x in &self.functions {
            x.add_monomorphs(self, &mut monomorphs);
        }

        // Insert the monomorphs into self
        for monomorph in monomorphs.drain_structs() {
            self.structs.try_insert(monomorph);
        }
        for monomorph in monomorphs.drain_unions() {
            self.unions.try_insert(monomorph);
        }
        for monomorph in monomorphs.drain_opaques() {
            self.opaque_items.try_insert(monomorph);
        }
        for monomorph in monomorphs.drain_typedefs() {
            self.typedefs.try_insert(monomorph);
        }

        // Remove structs and opaque items that are generic
        self.opaque_items.filter(|x| x.generic_params.len() > 0);
        self.structs.filter(|x| x.generic_params.len() > 0);
        self.unions.filter(|x| x.generic_params.len() > 0);
        self.typedefs.filter(|x| x.generic_params.len() > 0);

        // Mangle the paths that remain
        self.unions
            .for_all_items_mut(|x| x.mangle_paths(&monomorphs));
        self.structs
            .for_all_items_mut(|x| x.mangle_paths(&monomorphs));
        self.typedefs
            .for_all_items_mut(|x| x.mangle_paths(&monomorphs));
        for x in &mut self.functions {
            x.mangle_paths(&monomorphs);
        }
    }
}
<|endoftext|>


<commit_before>// (C)opyleft 2013-2019 Frank Denis

//! Bloom filter for Rust
//!
//! This is a simple but fast Bloom filter implementation, that requires only
//! 2 hash functions, generated with SipHash-1-3 using randomized keys.
//!

#![crate_name = "bloomfilter"]
#![crate_type = "rlib"]
#![warn(non_camel_case_types, non_upper_case_globals, unused_qualifications)]
#![allow(clippy::unreadable_literal, clippy::bool_comparison)]
extern crate bit_vec;
extern crate rand;
extern crate siphasher;

use bit_vec::BitVec;
use rand::prelude::*;
use siphasher::sip::SipHasher13;
use std::cmp;
use std::f64;
use std::hash::{Hash, Hasher};
use std::marker::PhantomData;

#[cfg(test)]
use rand::Rng;

/// Bloom filter structure
pub struct Bloom<T: ?Sized> {
    bitmap: BitVec,
    bitmap_bits: u64,
    k_num: u32,
    sips: [SipHasher13; 2],

    _phantom: PhantomData<T>,
}

impl<T: ?Sized> Bloom<T> {
    /// Create a new bloom filter structure.
    /// bitmap_size is the size in bytes (not bits) that will be allocated in memory
    /// items_count is an estimation of the maximum number of items to store.
    pub fn new(bitmap_size: usize, items_count: usize) -> Self {
        assert!(bitmap_size > 0 && items_count > 0);
        let bitmap_bits = (bitmap_size as u64) * 8u64;
        let k_num = Self::optimal_k_num(bitmap_bits, items_count);
        let bitmap = BitVec::from_elem(bitmap_bits as usize, false);
        let sips = [Self::sip_new(), Self::sip_new()];
        Self {
            bitmap,
            bitmap_bits,
            k_num,
            sips,
            _phantom: PhantomData,
        }
    }

    /// Create a new bloom filter structure.
    /// items_count is an estimation of the maximum number of items to store.
    /// fp_p is the wanted rate of false positives, in ]0.0, 1.0[
    pub fn new_for_fp_rate(items_count: usize, fp_p: f64) -> Self {
        let bitmap_size = Self::compute_bitmap_size(items_count, fp_p);
        Bloom::new(bitmap_size, items_count)
    }

    /// Create a bloom filter structure with an existing state.
    /// The state is assumed to be retrieved from an existing bloom filter.
    pub fn from_existing(
        bitmap: &[u8],
        bitmap_bits: u64,
        k_num: u32,
        sip_keys: [(u64, u64); 2],
    ) -> Self {
        let sips = [
            SipHasher13::new_with_keys(sip_keys[0].0, sip_keys[0].1),
            SipHasher13::new_with_keys(sip_keys[1].0, sip_keys[1].1),
        ];
        Self {
            bitmap: BitVec::from_bytes(bitmap),
            bitmap_bits,
            k_num,
            sips,
            _phantom: PhantomData,
        }
    }

    /// Compute a recommended bitmap size for items_count items
    /// and a fp_p rate of false positives.
    /// fp_p obviously has to be within the ]0.0, 1.0[ range.
    pub fn compute_bitmap_size(items_count: usize, fp_p: f64) -> usize {
        assert!(items_count > 0);
        assert!(fp_p > 0.0 && fp_p < 1.0);
        let log2 = f64::consts::LN_2;
        let log2_2 = log2 * log2;
        ((items_count as f64) * f64::ln(fp_p) / (-8.0 * log2_2)).ceil() as usize
    }

    /// Record the presence of an item.
    pub fn set(&mut self, item: &T)
    where
        T: Hash,
    {
        let mut hashes = [0u64, 0u64];
        for k_i in 0..self.k_num {
            let bit_offset = (self.bloom_hash(&mut hashes, &item, k_i) % self.bitmap_bits) as usize;
            self.bitmap.set(bit_offset, true);
        }
    }

    /// Check if an item is present in the set.
    /// There can be false positives, but no false negatives.
    pub fn check(&self, item: &T) -> bool
    where
        T: Hash,
    {
        let mut hashes = [0u64, 0u64];
        for k_i in 0..self.k_num {
            let bit_offset = (self.bloom_hash(&mut hashes, &item, k_i) % self.bitmap_bits) as usize;
            if self.bitmap.get(bit_offset).unwrap() == false {
                return false;
            }
        }
        true
    }

    /// Record the presence of an item in the set,
    /// and return the previous state of this item.
    pub fn check_and_set(&mut self, item: &T) -> bool
    where
        T: Hash,
    {
        let mut hashes = [0u64, 0u64];
        let mut found = true;
        for k_i in 0..self.k_num {
            let bit_offset = (self.bloom_hash(&mut hashes, &item, k_i) % self.bitmap_bits) as usize;
            if self.bitmap.get(bit_offset).unwrap() == false {
                found = false;
                self.bitmap.set(bit_offset, true);
            }
        }
        found
    }

    /// Return the bitmap as a vector of bytes
    pub fn bitmap(&self) -> Vec<u8> {
        self.bitmap.to_bytes()
    }

    /// Return the number of bits in the filter
    pub fn number_of_bits(&self) -> u64 {
        self.bitmap_bits
    }

    /// Return the number of hash functions used for `check` and `set`
    pub fn number_of_hash_functions(&self) -> u32 {
        self.k_num
    }

    /// Return the keys used by the sip hasher
    pub fn sip_keys(&self) -> [(u64, u64); 2] {
        [self.sips[0].keys(), self.sips[1].keys()]
    }

    fn optimal_k_num(bitmap_bits: u64, items_count: usize) -> u32 {
        let m = bitmap_bits as f64;
        let n = items_count as f64;
        let k_num = (m / n * f64::ln(2.0f64)).ceil() as u32;
        cmp::max(k_num, 1)
    }

    fn bloom_hash(&self, hashes: &mut [u64; 2], item: &T, k_i: u32) -> u64
    where
        T: Hash,
    {
        if k_i < 2 {
            let sip = &mut self.sips[k_i as usize].clone();
            item.hash(sip);
            let hash = sip.finish();
            hashes[k_i as usize] = hash;
            hash
        } else {
            ((hashes[0] as u128) + ((k_i as u128) * (hashes[1] as u128)) % 0xffffffffffffffc5)
                as u64
        }
    }

    /// Clear all of the bits in the filter, removing all keys from the set
    pub fn clear(&mut self) {
        self.bitmap.clear()
    }

    fn sip_new() -> SipHasher13 {
        let mut rng = thread_rng();
        SipHasher13::new_with_keys(rng.gen(), rng.gen())
    }
}

#[test]
fn bloom_test_set() {
    let mut rng = thread_rng();
    let mut bloom = Bloom::new(10, 80);
    let mut key = vec![0u8, 16];
    rng.fill_bytes(&mut key);
    assert!(bloom.check(&key) == false);
    bloom.set(&key);
    assert!(bloom.check(&key) == true);
}

#[test]
fn bloom_test_check_and_set() {
    let mut rng = thread_rng();
    let mut bloom = Bloom::new(10, 80);
    let mut key = vec![0u8, 16];
    rng.fill_bytes(&mut key);
    assert!(bloom.check_and_set(&key) == false);
    assert!(bloom.check_and_set(&key) == true);
}

#[test]
fn bloom_test_clear() {
    let mut rng = thread_rng();
    let mut bloom = Bloom::new(10, 80);
    let mut key = vec![0u8, 16];
    rng.fill_bytes(&mut key);
    bloom.set(&key);
    assert!(bloom.check(&key) == true);
    bloom.clear();
    assert!(bloom.check(&key) == false);
}

#[test]
fn bloom_test_load() {
    let mut rng = thread_rng();
    let mut original = Bloom::new(10, 80);
    let mut key = vec![0u8, 16];
    rng.fill_bytes(&mut key);
    original.set(&key);
    assert!(original.check(&key) == true);

    let cloned = Bloom::from_existing(
        &original.bitmap(),
        original.number_of_bits(),
        original.number_of_hash_functions(),
        original.sip_keys(),
    );
    assert!(cloned.check(&key) == true);
}
<commit_msg>Keep using wrapping operations (faster)<commit_after>// (C)opyleft 2013-2019 Frank Denis

//! Bloom filter for Rust
//!
//! This is a simple but fast Bloom filter implementation, that requires only
//! 2 hash functions, generated with SipHash-1-3 using randomized keys.
//!

#![crate_name = "bloomfilter"]
#![crate_type = "rlib"]
#![warn(non_camel_case_types, non_upper_case_globals, unused_qualifications)]
#![allow(clippy::unreadable_literal, clippy::bool_comparison)]
extern crate bit_vec;
extern crate rand;
extern crate siphasher;

use bit_vec::BitVec;
use rand::prelude::*;
use siphasher::sip::SipHasher13;
use std::cmp;
use std::f64;
use std::hash::{Hash, Hasher};
use std::marker::PhantomData;

#[cfg(test)]
use rand::Rng;

/// Bloom filter structure
pub struct Bloom<T: ?Sized> {
    bitmap: BitVec,
    bitmap_bits: u64,
    k_num: u32,
    sips: [SipHasher13; 2],

    _phantom: PhantomData<T>,
}

impl<T: ?Sized> Bloom<T> {
    /// Create a new bloom filter structure.
    /// bitmap_size is the size in bytes (not bits) that will be allocated in memory
    /// items_count is an estimation of the maximum number of items to store.
    pub fn new(bitmap_size: usize, items_count: usize) -> Self {
        assert!(bitmap_size > 0 && items_count > 0);
        let bitmap_bits = (bitmap_size as u64) * 8u64;
        let k_num = Self::optimal_k_num(bitmap_bits, items_count);
        let bitmap = BitVec::from_elem(bitmap_bits as usize, false);
        let sips = [Self::sip_new(), Self::sip_new()];
        Self {
            bitmap,
            bitmap_bits,
            k_num,
            sips,
            _phantom: PhantomData,
        }
    }

    /// Create a new bloom filter structure.
    /// items_count is an estimation of the maximum number of items to store.
    /// fp_p is the wanted rate of false positives, in ]0.0, 1.0[
    pub fn new_for_fp_rate(items_count: usize, fp_p: f64) -> Self {
        let bitmap_size = Self::compute_bitmap_size(items_count, fp_p);
        Bloom::new(bitmap_size, items_count)
    }

    /// Create a bloom filter structure with an existing state.
    /// The state is assumed to be retrieved from an existing bloom filter.
    pub fn from_existing(
        bitmap: &[u8],
        bitmap_bits: u64,
        k_num: u32,
        sip_keys: [(u64, u64); 2],
    ) -> Self {
        let sips = [
            SipHasher13::new_with_keys(sip_keys[0].0, sip_keys[0].1),
            SipHasher13::new_with_keys(sip_keys[1].0, sip_keys[1].1),
        ];
        Self {
            bitmap: BitVec::from_bytes(bitmap),
            bitmap_bits,
            k_num,
            sips,
            _phantom: PhantomData,
        }
    }

    /// Compute a recommended bitmap size for items_count items
    /// and a fp_p rate of false positives.
    /// fp_p obviously has to be within the ]0.0, 1.0[ range.
    pub fn compute_bitmap_size(items_count: usize, fp_p: f64) -> usize {
        assert!(items_count > 0);
        assert!(fp_p > 0.0 && fp_p < 1.0);
        let log2 = f64::consts::LN_2;
        let log2_2 = log2 * log2;
        ((items_count as f64) * f64::ln(fp_p) / (-8.0 * log2_2)).ceil() as usize
    }

    /// Record the presence of an item.
    pub fn set(&mut self, item: &T)
    where
        T: Hash,
    {
        let mut hashes = [0u64, 0u64];
        for k_i in 0..self.k_num {
            let bit_offset = (self.bloom_hash(&mut hashes, &item, k_i) % self.bitmap_bits) as usize;
            self.bitmap.set(bit_offset, true);
        }
    }

    /// Check if an item is present in the set.
    /// There can be false positives, but no false negatives.
    pub fn check(&self, item: &T) -> bool
    where
        T: Hash,
    {
        let mut hashes = [0u64, 0u64];
        for k_i in 0..self.k_num {
            let bit_offset = (self.bloom_hash(&mut hashes, &item, k_i) % self.bitmap_bits) as usize;
            if self.bitmap.get(bit_offset).unwrap() == false {
                return false;
            }
        }
        true
    }

    /// Record the presence of an item in the set,
    /// and return the previous state of this item.
    pub fn check_and_set(&mut self, item: &T) -> bool
    where
        T: Hash,
    {
        let mut hashes = [0u64, 0u64];
        let mut found = true;
        for k_i in 0..self.k_num {
            let bit_offset = (self.bloom_hash(&mut hashes, &item, k_i) % self.bitmap_bits) as usize;
            if self.bitmap.get(bit_offset).unwrap() == false {
                found = false;
                self.bitmap.set(bit_offset, true);
            }
        }
        found
    }

    /// Return the bitmap as a vector of bytes
    pub fn bitmap(&self) -> Vec<u8> {
        self.bitmap.to_bytes()
    }

    /// Return the number of bits in the filter
    pub fn number_of_bits(&self) -> u64 {
        self.bitmap_bits
    }

    /// Return the number of hash functions used for `check` and `set`
    pub fn number_of_hash_functions(&self) -> u32 {
        self.k_num
    }

    /// Return the keys used by the sip hasher
    pub fn sip_keys(&self) -> [(u64, u64); 2] {
        [self.sips[0].keys(), self.sips[1].keys()]
    }

    fn optimal_k_num(bitmap_bits: u64, items_count: usize) -> u32 {
        let m = bitmap_bits as f64;
        let n = items_count as f64;
        let k_num = (m / n * f64::ln(2.0f64)).ceil() as u32;
        cmp::max(k_num, 1)
    }

    fn bloom_hash(&self, hashes: &mut [u64; 2], item: &T, k_i: u32) -> u64
    where
        T: Hash,
    {
        if k_i < 2 {
            let sip = &mut self.sips[k_i as usize].clone();
            item.hash(sip);
            let hash = sip.finish();
            hashes[k_i as usize] = hash;
            hash
        } else {
            (hashes[0] as u128).wrapping_add((k_i as u128).wrapping_mul(hashes[1] as u128)) as u64
                % 0xffffffffffffffc5
        }
    }

    /// Clear all of the bits in the filter, removing all keys from the set
    pub fn clear(&mut self) {
        self.bitmap.clear()
    }

    fn sip_new() -> SipHasher13 {
        let mut rng = thread_rng();
        SipHasher13::new_with_keys(rng.gen(), rng.gen())
    }
}

#[test]
fn bloom_test_set() {
    let mut rng = thread_rng();
    let mut bloom = Bloom::new(10, 80);
    let mut key = vec![0u8, 16];
    rng.fill_bytes(&mut key);
    assert!(bloom.check(&key) == false);
    bloom.set(&key);
    assert!(bloom.check(&key) == true);
}

#[test]
fn bloom_test_check_and_set() {
    let mut rng = thread_rng();
    let mut bloom = Bloom::new(10, 80);
    let mut key = vec![0u8, 16];
    rng.fill_bytes(&mut key);
    assert!(bloom.check_and_set(&key) == false);
    assert!(bloom.check_and_set(&key) == true);
}

#[test]
fn bloom_test_clear() {
    let mut rng = thread_rng();
    let mut bloom = Bloom::new(10, 80);
    let mut key = vec![0u8, 16];
    rng.fill_bytes(&mut key);
    bloom.set(&key);
    assert!(bloom.check(&key) == true);
    bloom.clear();
    assert!(bloom.check(&key) == false);
}

#[test]
fn bloom_test_load() {
    let mut rng = thread_rng();
    let mut original = Bloom::new(10, 80);
    let mut key = vec![0u8, 16];
    rng.fill_bytes(&mut key);
    original.set(&key);
    assert!(original.check(&key) == true);

    let cloned = Bloom::from_existing(
        &original.bitmap(),
        original.number_of_bits(),
        original.number_of_hash_functions(),
        original.sip_keys(),
    );
    assert!(cloned.check(&key) == true);
}
<|endoftext|>


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

//! Implementation of `make clean` in rustbuild.
//!
//! Responsible for cleaning out a build directory of all old and stale
//! artifacts to prepare for a fresh build. Currently doesn't remove the
//! `build/cache` directory (download cache) or the `build/$target/llvm`
//! directory as we want that cached between builds.

use std::fs;
use std::io::{self, ErrorKind};
use std::path::Path;

use Build;

pub fn clean(build: &Build) {
    rm_rf(build, "tmp".as_ref());
    rm_rf(build, &build.out.join("tmp"));
    rm_rf(build, &build.out.join("dist"));

    for host in build.config.host.iter() {
        let entries = match build.out.join(host).read_dir() {
            Ok(iter) => iter,
            Err(_) => continue,
        };

        for entry in entries {
            let entry = t!(entry);
            if entry.file_name().to_str() == Some("llvm") {
                continue
            }
            let path = t!(entry.path().canonicalize());
            rm_rf(build, &path);
        }
    }
}

fn rm_rf(build: &Build, path: &Path) {
    if !path.exists() {
        return
    }
    if path.is_file() {
        return do_op(path, "remove file", |p| fs::remove_file(p));
    }

    for file in t!(fs::read_dir(path)) {
        let file = t!(file).path();

        if file.is_dir() {
            rm_rf(build, &file);
        } else {
            // On windows we can't remove a readonly file, and git will
            // often clone files as readonly. As a result, we have some
            // special logic to remove readonly files on windows.
            do_op(&file, "remove file", |p| fs::remove_file(p));
        }
    }
    do_op(path, "remove dir", |p| fs::remove_dir(p));
}

fn do_op<F>(path: &Path, desc: &str, mut f: F)
    where F: FnMut(&Path) -> io::Result<()>
{
    match f(path) {
        Ok(()) => {}
        Err(ref e) if cfg!(windows) &&
                      e.kind() == ErrorKind::PermissionDenied => {
            let mut p = t!(path.metadata()).permissions();
            p.set_readonly(false);
            t!(fs::set_permissions(path, p));
            f(path).unwrap_or_else(|e| {
                panic!("failed to {} {}: {}", desc, path.display(), e);
            })
        }
        Err(e) => {
            panic!("failed to {} {}: {}", desc, path.display(), e);
        }
    }
}
<commit_msg>Remove unused param from bootstrap::clean::rm_rf<commit_after>// Copyright 2016 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Implementation of `make clean` in rustbuild.
//!
//! Responsible for cleaning out a build directory of all old and stale
//! artifacts to prepare for a fresh build. Currently doesn't remove the
//! `build/cache` directory (download cache) or the `build/$target/llvm`
//! directory as we want that cached between builds.

use std::fs;
use std::io::{self, ErrorKind};
use std::path::Path;

use Build;

pub fn clean(build: &Build) {
    rm_rf("tmp".as_ref());
    rm_rf(&build.out.join("tmp"));
    rm_rf(&build.out.join("dist"));

    for host in build.config.host.iter() {
        let entries = match build.out.join(host).read_dir() {
            Ok(iter) => iter,
            Err(_) => continue,
        };

        for entry in entries {
            let entry = t!(entry);
            if entry.file_name().to_str() == Some("llvm") {
                continue
            }
            let path = t!(entry.path().canonicalize());
            rm_rf(&path);
        }
    }
}

fn rm_rf(path: &Path) {
    if !path.exists() {
        return
    }
    if path.is_file() {
        return do_op(path, "remove file", |p| fs::remove_file(p));
    }

    for file in t!(fs::read_dir(path)) {
        let file = t!(file).path();

        if file.is_dir() {
            rm_rf(&file);
        } else {
            // On windows we can't remove a readonly file, and git will
            // often clone files as readonly. As a result, we have some
            // special logic to remove readonly files on windows.
            do_op(&file, "remove file", |p| fs::remove_file(p));
        }
    }
    do_op(path, "remove dir", |p| fs::remove_dir(p));
}

fn do_op<F>(path: &Path, desc: &str, mut f: F)
    where F: FnMut(&Path) -> io::Result<()>
{
    match f(path) {
        Ok(()) => {}
        Err(ref e) if cfg!(windows) &&
                      e.kind() == ErrorKind::PermissionDenied => {
            let mut p = t!(path.metadata()).permissions();
            p.set_readonly(false);
            t!(fs::set_permissions(path, p));
            f(path).unwrap_or_else(|e| {
                panic!("failed to {} {}: {}", desc, path.display(), e);
            })
        }
        Err(e) => {
            panic!("failed to {} {}: {}", desc, path.display(), e);
        }
    }
}
<|endoftext|>


<commit_before>use quote::Tokens;
use syn::{Data, DeriveInput, Field, Fields};
use utils::{field_idents, get_field_types, named_to_vec, number_idents, unnamed_to_vec};
use quote::ToTokens;

/// Provides the hook to expand `#[derive(Constructor)]` into an implementation of `Constructor`
pub fn expand(input: &DeriveInput, _: &str) -> Tokens {
    let input_type = &input.ident;
    let field_vec: Vec<_>;
    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
    let (field_names, fields) = match input.data {
        Data::Struct(ref data_struct) => match data_struct.fields {
            Fields::Unnamed(ref fields) => {
                field_vec = unnamed_to_vec(fields);
                (tuple_field_names(&field_vec), field_vec)
            }
            Fields::Named(ref fields) => {
                field_vec = named_to_vec(fields);
                (struct_field_names(&field_vec), field_vec)
            }
            Fields::Unit => (vec![], vec![]),
        },
        _ => panic!("Only structs can derive a constructor"),
    };

    let original_types = &get_field_types(&fields);

    quote!{
        impl#impl_generics ::std::convert::From<#input_type#ty_generics> for
            (#(#original_types),*) #where_clause {

            #[allow(unused_variables)]
            #[inline]
            fn from(original: #input_type#ty_generics) -> (#(#original_types),*) {
                (#(original.#field_names),*)
            }
        }
    }
}

fn tuple_field_names(fields: &[&Field]) -> Vec<Tokens> {
    number_idents(fields.len())
        .iter()
        .map(|f| f.into_tokens())
        .collect()
}

fn struct_field_names(fields: &[&Field]) -> Vec<Tokens> {
    field_idents(fields)
        .iter()
        .map(|f| (*f).into_tokens())
        .collect()
}
<commit_msg>Fix leftovers from the copy from constructor.rs in into.rs<commit_after>use quote::Tokens;
use syn::{Data, DeriveInput, Field, Fields};
use utils::{field_idents, get_field_types, named_to_vec, number_idents, unnamed_to_vec};
use quote::ToTokens;

/// Provides the hook to expand `#[derive(Into)]` into an implementation of `Into`
pub fn expand(input: &DeriveInput, _: &str) -> Tokens {
    let input_type = &input.ident;
    let field_vec: Vec<_>;
    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
    let (field_names, fields) = match input.data {
        Data::Struct(ref data_struct) => match data_struct.fields {
            Fields::Unnamed(ref fields) => {
                field_vec = unnamed_to_vec(fields);
                (tuple_field_names(&field_vec), field_vec)
            }
            Fields::Named(ref fields) => {
                field_vec = named_to_vec(fields);
                (struct_field_names(&field_vec), field_vec)
            }
            Fields::Unit => (vec![], vec![]),
        },
        _ => panic!("Only structs can derive Into"),
    };

    let original_types = &get_field_types(&fields);

    quote!{
        impl#impl_generics ::std::convert::From<#input_type#ty_generics> for
            (#(#original_types),*) #where_clause {

            #[allow(unused_variables)]
            #[inline]
            fn from(original: #input_type#ty_generics) -> (#(#original_types),*) {
                (#(original.#field_names),*)
            }
        }
    }
}

fn tuple_field_names(fields: &[&Field]) -> Vec<Tokens> {
    number_idents(fields.len())
        .iter()
        .map(|f| f.into_tokens())
        .collect()
}

fn struct_field_names(fields: &[&Field]) -> Vec<Tokens> {
    field_idents(fields)
        .iter()
        .map(|f| (*f).into_tokens())
        .collect()
}
<|endoftext|>


<commit_before>use crate::progress_bar::ProgressBar;
use std::convert::TryFrom;
use std::io::{self, IoSliceMut};
use std::iter::FusedIterator;
#[cfg(feature = "tokio")]
use std::{
    pin::Pin,
    task::{Context, Poll},
};
#[cfg(feature = "tokio")]
use tokio::io::{ReadBuf, SeekFrom};

/// Wraps an iterator to display its progress.
pub trait ProgressIterator
where
    Self: Sized + Iterator,
{
    /// Wrap an iterator with default styling. Uses `Iterator::size_hint` to get length.
    /// Returns `Some(..)` only if `size_hint.1` is `Some`. If you want to create a progress bar
    /// even if `size_hint.1` returns `None` use `progress_count` or `progress_with` instead.
    fn try_progress(self) -> Option<ProgressBarIter<Self>> {
        self.size_hint()
            .1
            .map(|len| self.progress_count(u64::try_from(len).unwrap()))
    }

    /// Wrap an iterator with default styling.
    fn progress(self) -> ProgressBarIter<Self>
    where
        Self: ExactSizeIterator,
    {
        let len = u64::try_from(self.len()).unwrap();
        self.progress_count(len)
    }

    /// Wrap an iterator with an explicit element count.
    fn progress_count(self, len: u64) -> ProgressBarIter<Self> {
        self.progress_with(ProgressBar::new(len))
    }

    /// Wrap an iterator with a custom progress bar.
    fn progress_with(self, progress: ProgressBar) -> ProgressBarIter<Self>;

    /// Wrap an iterator with a progress bar and style it.
    fn progress_with_style(self, style: crate::ProgressStyle) -> ProgressBarIter<Self>
    where
        Self: ExactSizeIterator,
    {
        let len = u64::try_from(self.len()).unwrap();
        let bar = ProgressBar::new(len).with_style(style);
        self.progress_with(bar)
    }
}

/// Wraps an iterator to display its progress.
#[derive(Debug)]
pub struct ProgressBarIter<T> {
    pub(crate) it: T,
    pub progress: ProgressBar,
}

impl<S, T: Iterator<Item = S>> Iterator for ProgressBarIter<T> {
    type Item = S;

    fn next(&mut self) -> Option<Self::Item> {
        let item = self.it.next();

        if item.is_some() {
            self.progress.inc(1);
        } else if !self.progress.is_finished() {
            self.progress.finish_using_style();
        }

        item
    }
}

impl<T: ExactSizeIterator> ExactSizeIterator for ProgressBarIter<T> {
    fn len(&self) -> usize {
        self.it.len()
    }
}

impl<T: DoubleEndedIterator> DoubleEndedIterator for ProgressBarIter<T> {
    fn next_back(&mut self) -> Option<Self::Item> {
        let item = self.it.next_back();

        if item.is_some() {
            self.progress.inc(1);
        } else if !self.progress.is_finished() {
            self.progress.finish_using_style();
        }

        item
    }
}

impl<T: FusedIterator> FusedIterator for ProgressBarIter<T> {}

impl<R: io::Read> io::Read for ProgressBarIter<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let inc = self.it.read(buf)?;
        self.progress.inc(inc as u64);
        Ok(inc)
    }

    fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
        let inc = self.it.read_vectored(bufs)?;
        self.progress.inc(inc as u64);
        Ok(inc)
    }

    fn read_to_string(&mut self, buf: &mut String) -> io::Result<usize> {
        let inc = self.it.read_to_string(buf)?;
        self.progress.inc(inc as u64);
        Ok(inc)
    }

    fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
        self.it.read_exact(buf)?;
        self.progress.inc(buf.len() as u64);
        Ok(())
    }
}

impl<R: io::BufRead> io::BufRead for ProgressBarIter<R> {
    fn fill_buf(&mut self) -> io::Result<&[u8]> {
        self.it.fill_buf()
    }

    fn consume(&mut self, amt: usize) {
        self.it.consume(amt);
        self.progress.inc(amt as u64);
    }
}

impl<S: io::Seek> io::Seek for ProgressBarIter<S> {
    fn seek(&mut self, f: io::SeekFrom) -> io::Result<u64> {
        self.it.seek(f).map(|pos| {
            self.progress.set_position(pos);
            pos
        })
    }
    // Pass this through to preserve optimizations that the inner I/O object may use here
    // Also avoid sending a set_position update when the position hasn't changed
    fn stream_position(&mut self) -> io::Result<u64> {
        self.it.stream_position()
    }
}

#[cfg(feature = "tokio")]
impl<W: tokio::io::AsyncWrite + Unpin> tokio::io::AsyncWrite for ProgressBarIter<W> {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        Pin::new(&mut self.it).poll_write(cx, buf).map(|poll| {
            poll.map(|inc| {
                self.progress.inc(inc as u64);
                inc
            })
        })
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.it).poll_flush(cx)
    }

    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.it).poll_shutdown(cx)
    }
}

#[cfg(feature = "tokio")]
impl<W: tokio::io::AsyncRead + Unpin> tokio::io::AsyncRead for ProgressBarIter<W> {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        let prev_len = buf.filled().len() as u64;
        if let Poll::Ready(e) = Pin::new(&mut self.it).poll_read(cx, buf) {
            self.progress.inc(buf.filled().len() as u64 - prev_len);
            Poll::Ready(e)
        } else {
            Poll::Pending
        }
    }
}

#[cfg(feature = "tokio")]
impl<W: tokio::io::AsyncSeek + Unpin> tokio::io::AsyncSeek for ProgressBarIter<W> {
    fn start_seek(mut self: Pin<&mut Self>, position: SeekFrom) -> io::Result<()> {
        Pin::new(&mut self.it).start_seek(position)
    }

    fn poll_complete(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<u64>> {
        Pin::new(&mut self.it).poll_complete(cx)
    }
}

#[cfg(feature = "tokio")]
impl<W: tokio::io::AsyncBufRead + Unpin + tokio::io::AsyncRead> tokio::io::AsyncBufRead
    for ProgressBarIter<W>
{
    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
        let this = self.get_mut();
        let result = Pin::new(&mut this.it).poll_fill_buf(cx);
        if let Poll::Ready(Ok(buf)) = &result {
            this.progress.inc(buf.len() as u64);
        }
        result
    }

    fn consume(mut self: Pin<&mut Self>, amt: usize) {
        Pin::new(&mut self.it).consume(amt);
    }
}

impl<W: io::Write> io::Write for ProgressBarIter<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.it.write(buf).map(|inc| {
            self.progress.inc(inc as u64);
            inc
        })
    }

    fn write_vectored(&mut self, bufs: &[io::IoSlice]) -> io::Result<usize> {
        self.it.write_vectored(bufs).map(|inc| {
            self.progress.inc(inc as u64);
            inc
        })
    }

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

    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
        self.it.write_all(buf).map(|()| {
            self.progress.inc(buf.len() as u64);
        })
    }

    // write_fmt can not be captured with reasonable effort.
    // as it uses write_all internally by default that should not be a problem.
    // fn write_fmt(&mut self, fmt: fmt::Arguments) -> io::Result<()>;
}

impl<S, T: Iterator<Item = S>> ProgressIterator for T {
    fn progress_with(self, progress: ProgressBar) -> ProgressBarIter<Self> {
        ProgressBarIter { it: self, progress }
    }
}

#[cfg(test)]
mod test {
    use crate::iter::{ProgressBarIter, ProgressIterator};
    use crate::progress_bar::ProgressBar;
    use crate::ProgressStyle;

    #[test]
    fn it_can_wrap_an_iterator() {
        let v = vec![1, 2, 3];
        let wrap = |it: ProgressBarIter<_>| {
            assert_eq!(it.map(|x| x * 2).collect::<Vec<_>>(), vec![2, 4, 6]);
        };

        wrap(v.iter().progress());
        wrap(v.iter().progress_count(3));
        wrap({
            let pb = ProgressBar::new(v.len() as u64);
            v.iter().progress_with(pb)
        });
        wrap({
            let style = ProgressStyle::default_bar().template("{wide_bar:.red} {percent}/100%");
            v.iter().progress_with_style(style)
        });
    }
}
<commit_msg>Add builder methods for ProgressBarIter<commit_after>use crate::progress_bar::ProgressBar;
use crate::style::ProgressStyle;
use std::borrow::Cow;
use std::convert::TryFrom;
use std::io::{self, IoSliceMut};
use std::iter::FusedIterator;
use std::time::Duration;
#[cfg(feature = "tokio")]
use std::{
    pin::Pin,
    task::{Context, Poll},
};
#[cfg(feature = "tokio")]
use tokio::io::{ReadBuf, SeekFrom};

/// Wraps an iterator to display its progress.
pub trait ProgressIterator
where
    Self: Sized + Iterator,
{
    /// Wrap an iterator with default styling. Uses `Iterator::size_hint` to get length.
    /// Returns `Some(..)` only if `size_hint.1` is `Some`. If you want to create a progress bar
    /// even if `size_hint.1` returns `None` use `progress_count` or `progress_with` instead.
    fn try_progress(self) -> Option<ProgressBarIter<Self>> {
        self.size_hint()
            .1
            .map(|len| self.progress_count(u64::try_from(len).unwrap()))
    }

    /// Wrap an iterator with default styling.
    fn progress(self) -> ProgressBarIter<Self>
    where
        Self: ExactSizeIterator,
    {
        let len = u64::try_from(self.len()).unwrap();
        self.progress_count(len)
    }

    /// Wrap an iterator with an explicit element count.
    fn progress_count(self, len: u64) -> ProgressBarIter<Self> {
        self.progress_with(ProgressBar::new(len))
    }

    /// Wrap an iterator with a custom progress bar.
    fn progress_with(self, progress: ProgressBar) -> ProgressBarIter<Self>;

    /// Wrap an iterator with a progress bar and style it.
    fn progress_with_style(self, style: crate::ProgressStyle) -> ProgressBarIter<Self>
    where
        Self: ExactSizeIterator,
    {
        let len = u64::try_from(self.len()).unwrap();
        let bar = ProgressBar::new(len).with_style(style);
        self.progress_with(bar)
    }
}

/// Wraps an iterator to display its progress.
#[derive(Debug)]
pub struct ProgressBarIter<T> {
    pub(crate) it: T,
    pub progress: ProgressBar,
}

impl<T> ProgressBarIter<T> {
    /// Builder-like function for setting underlying progress bar's style.
    ///
    /// See [ProgressBar::with_style].
    pub fn with_style(mut self, style: ProgressStyle) -> ProgressBarIter<T> {
        self.progress = self.progress.with_style(style);
        self
    }

    /// Builder-like function for setting underlying progress bar's prefix.
    ///
    /// See [ProgressBar::with_prefix].
    pub fn with_prefix(mut self, prefix: impl Into<Cow<'static, str>>) -> ProgressBarIter<T> {
        self.progress = self.progress.with_prefix(prefix);
        self
    }

    /// Builder-like function for setting underlying progress bar's message.
    ///
    /// See [ProgressBar::with_message].
    pub fn with_message(mut self, message: impl Into<Cow<'static, str>>) -> ProgressBarIter<T> {
        self.progress = self.progress.with_message(message);
        self
    }

    /// Builder-like function for setting underlying progress bar's position.
    ///
    /// See [ProgressBar::with_position].
    pub fn with_position(mut self, position: u64) -> ProgressBarIter<T> {
        self.progress = self.progress.with_position(position);
        self
    }

    /// Builder-like function for setting underlying progress bar's elapsed time.
    ///
    /// See [ProgressBar::with_elapsed].
    pub fn with_elapsed(mut self, elapsed: Duration) -> ProgressBarIter<T> {
        self.progress = self.progress.with_elapsed(elapsed);
        self
    }
}

impl<S, T: Iterator<Item = S>> Iterator for ProgressBarIter<T> {
    type Item = S;

    fn next(&mut self) -> Option<Self::Item> {
        let item = self.it.next();

        if item.is_some() {
            self.progress.inc(1);
        } else if !self.progress.is_finished() {
            self.progress.finish_using_style();
        }

        item
    }
}

impl<T: ExactSizeIterator> ExactSizeIterator for ProgressBarIter<T> {
    fn len(&self) -> usize {
        self.it.len()
    }
}

impl<T: DoubleEndedIterator> DoubleEndedIterator for ProgressBarIter<T> {
    fn next_back(&mut self) -> Option<Self::Item> {
        let item = self.it.next_back();

        if item.is_some() {
            self.progress.inc(1);
        } else if !self.progress.is_finished() {
            self.progress.finish_using_style();
        }

        item
    }
}

impl<T: FusedIterator> FusedIterator for ProgressBarIter<T> {}

impl<R: io::Read> io::Read for ProgressBarIter<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let inc = self.it.read(buf)?;
        self.progress.inc(inc as u64);
        Ok(inc)
    }

    fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
        let inc = self.it.read_vectored(bufs)?;
        self.progress.inc(inc as u64);
        Ok(inc)
    }

    fn read_to_string(&mut self, buf: &mut String) -> io::Result<usize> {
        let inc = self.it.read_to_string(buf)?;
        self.progress.inc(inc as u64);
        Ok(inc)
    }

    fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
        self.it.read_exact(buf)?;
        self.progress.inc(buf.len() as u64);
        Ok(())
    }
}

impl<R: io::BufRead> io::BufRead for ProgressBarIter<R> {
    fn fill_buf(&mut self) -> io::Result<&[u8]> {
        self.it.fill_buf()
    }

    fn consume(&mut self, amt: usize) {
        self.it.consume(amt);
        self.progress.inc(amt as u64);
    }
}

impl<S: io::Seek> io::Seek for ProgressBarIter<S> {
    fn seek(&mut self, f: io::SeekFrom) -> io::Result<u64> {
        self.it.seek(f).map(|pos| {
            self.progress.set_position(pos);
            pos
        })
    }
    // Pass this through to preserve optimizations that the inner I/O object may use here
    // Also avoid sending a set_position update when the position hasn't changed
    fn stream_position(&mut self) -> io::Result<u64> {
        self.it.stream_position()
    }
}

#[cfg(feature = "tokio")]
impl<W: tokio::io::AsyncWrite + Unpin> tokio::io::AsyncWrite for ProgressBarIter<W> {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        Pin::new(&mut self.it).poll_write(cx, buf).map(|poll| {
            poll.map(|inc| {
                self.progress.inc(inc as u64);
                inc
            })
        })
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.it).poll_flush(cx)
    }

    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.it).poll_shutdown(cx)
    }
}

#[cfg(feature = "tokio")]
impl<W: tokio::io::AsyncRead + Unpin> tokio::io::AsyncRead for ProgressBarIter<W> {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        let prev_len = buf.filled().len() as u64;
        if let Poll::Ready(e) = Pin::new(&mut self.it).poll_read(cx, buf) {
            self.progress.inc(buf.filled().len() as u64 - prev_len);
            Poll::Ready(e)
        } else {
            Poll::Pending
        }
    }
}

#[cfg(feature = "tokio")]
impl<W: tokio::io::AsyncSeek + Unpin> tokio::io::AsyncSeek for ProgressBarIter<W> {
    fn start_seek(mut self: Pin<&mut Self>, position: SeekFrom) -> io::Result<()> {
        Pin::new(&mut self.it).start_seek(position)
    }

    fn poll_complete(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<u64>> {
        Pin::new(&mut self.it).poll_complete(cx)
    }
}

#[cfg(feature = "tokio")]
impl<W: tokio::io::AsyncBufRead + Unpin + tokio::io::AsyncRead> tokio::io::AsyncBufRead
    for ProgressBarIter<W>
{
    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
        let this = self.get_mut();
        let result = Pin::new(&mut this.it).poll_fill_buf(cx);
        if let Poll::Ready(Ok(buf)) = &result {
            this.progress.inc(buf.len() as u64);
        }
        result
    }

    fn consume(mut self: Pin<&mut Self>, amt: usize) {
        Pin::new(&mut self.it).consume(amt);
    }
}

impl<W: io::Write> io::Write for ProgressBarIter<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.it.write(buf).map(|inc| {
            self.progress.inc(inc as u64);
            inc
        })
    }

    fn write_vectored(&mut self, bufs: &[io::IoSlice]) -> io::Result<usize> {
        self.it.write_vectored(bufs).map(|inc| {
            self.progress.inc(inc as u64);
            inc
        })
    }

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

    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
        self.it.write_all(buf).map(|()| {
            self.progress.inc(buf.len() as u64);
        })
    }

    // write_fmt can not be captured with reasonable effort.
    // as it uses write_all internally by default that should not be a problem.
    // fn write_fmt(&mut self, fmt: fmt::Arguments) -> io::Result<()>;
}

impl<S, T: Iterator<Item = S>> ProgressIterator for T {
    fn progress_with(self, progress: ProgressBar) -> ProgressBarIter<Self> {
        ProgressBarIter { it: self, progress }
    }
}

#[cfg(test)]
mod test {
    use crate::iter::{ProgressBarIter, ProgressIterator};
    use crate::progress_bar::ProgressBar;
    use crate::ProgressStyle;

    #[test]
    fn it_can_wrap_an_iterator() {
        let v = vec![1, 2, 3];
        let wrap = |it: ProgressBarIter<_>| {
            assert_eq!(it.map(|x| x * 2).collect::<Vec<_>>(), vec![2, 4, 6]);
        };

        wrap(v.iter().progress());
        wrap(v.iter().progress_count(3));
        wrap({
            let pb = ProgressBar::new(v.len() as u64);
            v.iter().progress_with(pb)
        });
        wrap({
            let style = ProgressStyle::default_bar().template("{wide_bar:.red} {percent}/100%");
            v.iter().progress_with_style(style)
        });
    }
}
<|endoftext|>


<commit_before>// An Iterator adapter that walks through all the elements in the Iterator,
// converts them to Strings and joins them to one big String, seperated by
// some seperator string slice.
trait StringJoiner {
  fn join(&mut self, sep: &str) -> String;
}

// Implement it for all Iterators with Elements convertable into a String
impl<A: ToString, T: Iterator<A>> StringJoiner for T {
  fn join(&mut self, sep: &str) -> String {
    match self.next() {
      Some(elem) => {
        let mut output = elem.to_string();
        for elem in *self {
          output.push_str(sep);
          output.push_str(elem.to_string().as_slice());
        }
        output
      }
      None => String::new()
    }
  }
}

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

  assert_eq!(many.iter().join(", ").as_slice(), "1, 2, 3");
  assert_eq!( one.iter().join(", ").as_slice(), "1");
  assert_eq!(none.iter().join(", ").as_slice(), "");
}
<commit_msg>made StringJoiner pub and added docs<commit_after>// An Iterator adapter that walks through all the elements in the Iterator,
// converts them to Strings and joins them to one big String, seperated by
// some seperator string slice.
pub trait StringJoiner {
  /// Converts all elements the Iterator yields to Strings,
  /// then combines them all into one String, seperated by sep.
  ///
  /// # Example
  ///
  /// ```
  /// use collect::iter::StringJoiner;
  ///
  /// let vec = vec![1u,2,3];
  /// assert_eq!(vec.iter().join(", ").as_slice(), "1, 2, 3");
  /// ```
  fn join(&mut self, sep: &str) -> String;
}

// Implement it for all Iterators with Elements convertable into a String
impl<A: ToString, T: Iterator<A>> StringJoiner for T {
  fn join(&mut self, sep: &str) -> String {
    match self.next() {
      Some(elem) => {
        let mut output = elem.to_string();
        for elem in *self {
          output.push_str(sep);
          output.push_str(elem.to_string().as_slice());
        }
        output
      }
      None => String::new()
    }
  }
}

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

  assert_eq!(many.iter().join(", ").as_slice(), "1, 2, 3");
  assert_eq!( one.iter().join(", ").as_slice(), "1");
  assert_eq!(none.iter().join(", ").as_slice(), "");
}
<|endoftext|>