<commit_before>extern crate criterion;

use criterion::{Bencher, Criterion, ParameterizedBenchmark, criterion_group, criterion_main};
use image::{ColorType, bmp::BMPEncoder};

use std::fs::File;
use std::io::BufWriter;

fn encode_gray_test(criterion: &mut Criterion) {
    let counts = vec![100usize, 200, 400, 800, 1200];
    fn raw_vec(b: &mut Bencher, s: &usize) {
        let s = *s;
        let im = vec![0; s * s];
        let mut v = vec![0; s * s / 4 + 4096 + s * 4];
        b.iter(|| {
            v.clear();
            let mut x = BMPEncoder::new(&mut v);
            x.encode(&im, (s / 2) as u32, (s / 2) as u32, ColorType::L8).unwrap();
        }
        )
    }

    fn buf_vec(b: &mut Bencher, s: &usize) {
        let s = *s;
        let im = vec![0; s * s];
        let mut v = vec![0; s * s / 4 + 4096 + s * 4];
        b.iter(|| {
            v.clear();
            let mut buf = BufWriter::new(&mut v);
            let mut x = BMPEncoder::new(&mut buf);
            x.encode(&im, (s / 2) as u32, (s / 2) as u32, ColorType::L8).unwrap();
        }
        )
    }

    fn buf_file(b: &mut Bencher, s: &usize) {
        let s = *s;
        let im = vec![0; s * s];
        b.iter(|| {
            let mut f = File::open("temp.bmp").unwrap();
            let mut buf = BufWriter::new(&mut f);
            let mut x = BMPEncoder::new(&mut buf);
            x.encode(&im, (s / 2) as u32, (s / 2) as u32, ColorType::L8).unwrap();
        }
        )
    }

    criterion.bench("encode_gray",
                    ParameterizedBenchmark::new("raw_vec", raw_vec, counts)
                        .with_function("buf_vec", buf_vec)
                        .with_function("buf_file", buf_file));
}

fn encode_rgb_test(criterion: &mut Criterion) {
    let counts = vec![100usize, 200, 400, 800, 1200];
    fn raw_vec(b: &mut Bencher, s: &usize) {
        let s = *s;
        let im = vec![0; s * s];
        let mut v = vec![0; s * s / 4 + 4096 + s * 4];
        b.iter(|| {
            v.clear();
            let mut x = BMPEncoder::new(&mut v);
            x.encode(&im, (s / 2) as u32, (s / 2) as u32, ColorType::Rgb8).unwrap();
        }
        )
    }

    fn buf_vec(b: &mut Bencher, s: &usize) {
        let s = *s;
        let im = vec![0; s * s];
        let mut v = vec![0; s * s / 4 + 4096 + s * 4];
        b.iter(|| {
            v.clear();
            let mut buf = BufWriter::new(&mut v);
            let mut x = BMPEncoder::new(&mut buf);
            x.encode(&im, (s / 2) as u32, (s / 2) as u32, ColorType::Rgb8).unwrap();
        }
        )
    }

    fn buf_file(b: &mut Bencher, s: &usize) {
        let s = *s;
        let im = vec![0; s * s];
        b.iter(|| {
            let mut f = File::open("temp.bmp").unwrap();
            let mut buf = BufWriter::new(&mut f);
            let mut x = BMPEncoder::new(&mut buf);
            x.encode(&im, (s / 2) as u32, (s / 2) as u32, ColorType::Rgb8).unwrap();
        }
        )
    }

    criterion.bench("encode_rgb",
                    ParameterizedBenchmark::new("raw_vec", raw_vec, counts)
                        .with_function("buf_vec", buf_vec)
                        .with_function("buf_file", buf_file));
}

fn encode_rgba_test(criterion: &mut Criterion) {
    let counts = vec![100usize, 200, 400, 800, 1200];
    fn raw_vec(b: &mut Bencher, s: &usize) {
        let s = *s;
        let im = vec![0; s * s];
        let mut v = vec![0; s * s / 4 + 4096 + s * 4];
        b.iter(|| {
            v.clear();
            let mut x = BMPEncoder::new(&mut v);
            x.encode(&im, (s / 2) as u32, (s / 2) as u32, ColorType::Rgba8).unwrap();
        }
        )
    }

    fn buf_vec(b: &mut Bencher, s: &usize) {
        let s = *s;
        let im = vec![0; s * s];
        let mut v = vec![0; s * s / 4 + 4096 + s * 4];
        b.iter(|| {
            v.clear();
            let mut buf = BufWriter::new(&mut v);
            let mut x = BMPEncoder::new(&mut buf);
            x.encode(&im, (s / 2) as u32, (s / 2) as u32, ColorType::Rgba8).unwrap();
        }
        )
    }

    fn buf_file(b: &mut Bencher, s: &usize) {
        let s = *s;
        let im = vec![0; s * s];
        b.iter(|| {
            let mut f = File::open("temp.bmp").unwrap();
            let mut buf = BufWriter::new(&mut f);
            let mut x = BMPEncoder::new(&mut buf);
            x.encode(&im, (s / 2) as u32, (s / 2) as u32, ColorType::Rgba8).unwrap();
        }
        )
    }

    criterion.bench("encode_rgba",
                    ParameterizedBenchmark::new("raw_vec", raw_vec, counts)
                        .with_function("buf_vec", buf_vec)
                        .with_function("buf_file", buf_file));
}

criterion_group!(benches, encode_gray_test);
criterion_main!(benches);

<commit_msg>Update bmp encoding benchmark to criterion 0.3<commit_after>extern crate criterion;

use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
use image::{ColorType, bmp::BMPEncoder};

use std::fs::File;
use std::io::{BufWriter, Seek, SeekFrom};

fn encode_gray_test(criterion: &mut Criterion) {
    let counts = vec![100u32, 200, 400];

    let mut group = criterion.benchmark_group("encode_gray");
    for size in counts.iter().copied() {
        let bytes = size as usize;
        let im = vec![0; bytes * bytes];
        group.bench_with_input(BenchmarkId::new("rawvec", size), &im, |b, image| {
            let mut v = vec![];
            encode_raw(&mut v, &im, size, ColorType::L8);
            b.iter(|| encode_raw(&mut v, image, size, ColorType::L8));
        });
        group.bench_with_input(BenchmarkId::new("bufvec", size), &im, |b, image| {
            let mut v = vec![];
            encode_raw(&mut v, &im, size, ColorType::L8);
            b.iter(|| encode_bufvec(&mut v, image, size, ColorType::L8));
        });
        group.bench_with_input(BenchmarkId::new("file", size), &im, |b, image| {
            let file = File::create("temp.bmp").unwrap();
            b.iter(|| encode_file(&file, image, size, ColorType::L8));
        });
    }
}

criterion_group!(benches, encode_gray_test);
criterion_main!(benches);

fn encode_raw(into: &mut Vec<u8>, im: &[u8], size: u32, color: ColorType) {
    into.clear();
    let mut x = BMPEncoder::new(into);
    x.encode(im, size, size, color).unwrap();
}

fn encode_bufvec(into: &mut Vec<u8>, im: &[u8], size: u32, color: ColorType) {
    into.clear();
    let mut buf = BufWriter::new(into);
    let mut x = BMPEncoder::new(&mut buf);
    x.encode(im, size, size, color).unwrap();
}

fn encode_file(mut file: &File, im: &[u8], size: u32, color: ColorType) {
    file.seek(SeekFrom::Start(0)).unwrap();
    let mut buf = BufWriter::new(file);
    let mut x = BMPEncoder::new(&mut buf);
    x.encode(im, size, size, color).unwrap();
}
<|endoftext|>


<commit_before>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    ChainInfo, FullNode, HealthCheckError, LocalNode, LocalVersion, Node, Swarm, SwarmExt,
    Validator, Version,
};
use anyhow::{anyhow, Result};
use diem_config::config::NodeConfig;
use diem_genesis_tool::{fullnode_builder::FullnodeConfig, validator_builder::ValidatorBuilder};
use diem_sdk::{
    crypto::ed25519::Ed25519PrivateKey,
    types::{
        chain_id::ChainId, transaction::Transaction, waypoint::Waypoint, AccountKey, LocalAccount,
        PeerId,
    },
};
use std::{
    collections::HashMap,
    convert::TryFrom,
    fs, mem,
    num::NonZeroUsize,
    ops,
    path::{Path, PathBuf},
    sync::Arc,
    time::{Duration, Instant},
};
use tempfile::TempDir;

#[derive(Debug)]
pub enum SwarmDirectory {
    Persistent(PathBuf),
    Temporary(TempDir),
}

impl SwarmDirectory {
    pub fn persist(&mut self) {
        match self {
            SwarmDirectory::Persistent(_) => {}
            SwarmDirectory::Temporary(_) => {
                let mut temp = SwarmDirectory::Persistent(PathBuf::new());
                mem::swap(self, &mut temp);
                let _ = mem::replace(self, temp.into_persistent());
            }
        }
    }

    pub fn into_persistent(self) -> Self {
        match self {
            SwarmDirectory::Temporary(tempdir) => SwarmDirectory::Persistent(tempdir.into_path()),
            SwarmDirectory::Persistent(dir) => SwarmDirectory::Persistent(dir),
        }
    }
}

impl ops::Deref for SwarmDirectory {
    type Target = Path;

    fn deref(&self) -> &Self::Target {
        match self {
            SwarmDirectory::Persistent(dir) => dir.deref(),
            SwarmDirectory::Temporary(dir) => dir.path(),
        }
    }
}

impl AsRef<Path> for SwarmDirectory {
    fn as_ref(&self) -> &Path {
        match self {
            SwarmDirectory::Persistent(dir) => dir.as_ref(),
            SwarmDirectory::Temporary(dir) => dir.as_ref(),
        }
    }
}

#[derive(Debug)]
struct ValidatorNetworkAddressEncryptionKey {
    key: diem_sdk::types::network_address::encrypted::Key,
    version: diem_sdk::types::network_address::encrypted::KeyVersion,
}

pub struct LocalSwarmBuilder {
    versions: Arc<HashMap<Version, LocalVersion>>,
    initial_version: Option<Version>,
    template: NodeConfig,
    number_of_validators: NonZeroUsize,
    dir: Option<PathBuf>,
}

impl LocalSwarmBuilder {
    pub fn new(versions: Arc<HashMap<Version, LocalVersion>>) -> Self {
        Self {
            versions,
            initial_version: None,
            template: NodeConfig::default_for_validator(),
            number_of_validators: NonZeroUsize::new(1).unwrap(),
            dir: None,
        }
    }

    pub fn initial_version(mut self, initial_version: Version) -> Self {
        self.initial_version = Some(initial_version);
        self
    }

    pub fn template(mut self, template: NodeConfig) -> Self {
        self.template = template;
        self
    }

    pub fn number_of_validators(mut self, number_of_validators: NonZeroUsize) -> Self {
        self.number_of_validators = number_of_validators;
        self
    }

    pub fn dir<T: AsRef<Path>>(mut self, dir: T) -> Self {
        self.dir = Some(dir.as_ref().into());
        self
    }

    pub fn build<R>(self, rng: R) -> Result<LocalSwarm>
    where
        R: ::rand::RngCore + ::rand::CryptoRng,
    {
        let dir = if let Some(dir) = self.dir {
            if dir.exists() {
                fs::remove_dir_all(&dir)?;
            }
            fs::create_dir_all(&dir)?;
            SwarmDirectory::Persistent(dir)
        } else {
            SwarmDirectory::Temporary(TempDir::new()?)
        };

        let (root_keys, genesis, genesis_waypoint, validators) = ValidatorBuilder::new(
            &dir,
            diem_framework_releases::current_module_blobs().to_vec(),
        )
        .num_validators(self.number_of_validators)
        .template(self.template)
        .build(rng)?;

        // Get the initial version to start the nodes with, either the one provided or fallback to
        // using the the latest version
        let versions = self.versions;
        let initial_version = self.initial_version.unwrap_or_else(|| {
            versions
                .iter()
                .max_by(|v1, v2| v1.0.cmp(v2.0))
                .unwrap()
                .0
                .clone()
        });
        let version = versions.get(&initial_version).unwrap();

        let validators = validators
            .into_iter()
            .map(|v| {
                let node = LocalNode::new(version.to_owned(), v.name, v.directory)?;
                Ok((node.peer_id(), node))
            })
            .collect::<Result<HashMap<_, _>>>()?;

        let validator_network_address_encryption_key = ValidatorNetworkAddressEncryptionKey {
            key: root_keys.validator_network_address_encryption_key,
            version: root_keys.validator_network_address_encryption_key_version,
        };

        let root_account = LocalAccount::new(
            diem_sdk::types::account_config::diem_root_address(),
            AccountKey::from_private_key(root_keys.root_key),
            1,
        );

        // Designated dealer account is built to have the same private key as the TC account
        let designated_dealer_account = LocalAccount::new(
            diem_sdk::types::account_config::testnet_dd_account_address(),
            AccountKey::from_private_key(
                Ed25519PrivateKey::try_from(root_keys.treasury_compliance_key.to_bytes().as_ref())
                    .unwrap(),
            ),
            0,
        );

        let treasury_compliance_account = LocalAccount::new(
            diem_sdk::types::account_config::treasury_compliance_account_address(),
            AccountKey::from_private_key(root_keys.treasury_compliance_key),
            0,
        );

        Ok(LocalSwarm {
            node_name_counter: validators.len() as u64,
            genesis,
            genesis_waypoint,
            versions,
            validators,
            fullnodes: HashMap::new(),
            dir,
            validator_network_address_encryption_key,
            root_account,
            treasury_compliance_account,
            designated_dealer_account,
            chain_id: ChainId::test(),
        })
    }
}

#[derive(Debug)]
pub struct LocalSwarm {
    node_name_counter: u64,
    genesis: Transaction,
    genesis_waypoint: Waypoint,
    versions: Arc<HashMap<Version, LocalVersion>>,
    validators: HashMap<PeerId, LocalNode>,
    fullnodes: HashMap<PeerId, LocalNode>,
    dir: SwarmDirectory,
    validator_network_address_encryption_key: ValidatorNetworkAddressEncryptionKey,
    root_account: LocalAccount,
    treasury_compliance_account: LocalAccount,
    designated_dealer_account: LocalAccount,
    chain_id: ChainId,
}

impl LocalSwarm {
    pub fn builder(versions: Arc<HashMap<Version, LocalVersion>>) -> LocalSwarmBuilder {
        LocalSwarmBuilder::new(versions)
    }

    pub fn launch(&mut self) -> Result<()> {
        // Start all the validators
        for validator in self.validators.values_mut() {
            validator.start()?;
        }

        // Wait for all of them to startup
        let deadline = Instant::now() + Duration::from_secs(60);
        self.wait_for_startup()?;
        self.wait_for_connectivity(deadline)?;
        self.liveness_check(deadline)?;

        Ok(())
    }

    fn wait_for_startup(&mut self) -> Result<()> {
        let num_attempts = 10;
        let mut done = vec![false; self.validators.len()];
        for i in 0..num_attempts {
            println!("Wait for startup attempt: {} of {}", i, num_attempts);
            for (node, done) in self.validators.values_mut().zip(done.iter_mut()) {
                if *done {
                    continue;
                }
                match node.health_check() {
                    Ok(()) => *done = true,

                    Err(HealthCheckError::Unknown(e)) => {
                        return Err(anyhow!("Diem node '{}' is not running: {}", node.name(), e));
                    }
                    Err(HealthCheckError::NotRunning) => {
                        return Err(anyhow!("Diem node '{}' is not running", node.name()));
                    }
                    Err(HealthCheckError::RpcFailure(e)) => {
                        println!("rpc error: {}", e);
                        continue;
                    }
                }
            }

            // Check if all the nodes have been successfully launched
            if done.iter().all(|status| *status) {
                return Ok(());
            }

            ::std::thread::sleep(::std::time::Duration::from_millis(1000));
        }

        Err(anyhow!("Launching Swarm timed out"))
    }

    fn add_fullnode(&mut self, version: &Version, template: NodeConfig) -> Result<PeerId> {
        let name = self.node_name_counter.to_string();
        self.node_name_counter += 1;
        let fullnode_config = FullnodeConfig::public_fullnode(
            name,
            self.dir.as_ref(),
            template,
            &self.genesis_waypoint,
            &self.genesis,
        )?;

        let version = self.versions.get(version).unwrap();
        let mut fullnode = LocalNode::new(
            version.to_owned(),
            fullnode_config.name,
            fullnode_config.directory,
        )?;

        let peer_id = fullnode.peer_id();
        fullnode.start()?;

        self.fullnodes.insert(peer_id, fullnode);

        Ok(peer_id)
    }

    pub fn chain_id(&self) -> ChainId {
        self.chain_id
    }

    pub fn validators(&self) -> impl Iterator<Item = &LocalNode> {
        self.validators.values()
    }

    pub fn validators_mut(&mut self) -> impl Iterator<Item = &mut LocalNode> {
        self.validators.values_mut()
    }

    pub fn dir(&self) -> &Path {
        self.dir.as_ref()
    }
}

impl Swarm for LocalSwarm {
    fn health_check(&mut self) -> Result<()> {
        todo!()
    }

    fn validators<'a>(&'a self) -> Box<dyn Iterator<Item = &'a dyn Validator> + 'a> {
        Box::new(self.validators.values().map(|v| v as &'a dyn Validator))
    }

    fn validators_mut<'a>(&'a mut self) -> Box<dyn Iterator<Item = &'a mut dyn Validator> + 'a> {
        Box::new(
            self.validators
                .values_mut()
                .map(|v| v as &'a mut dyn Validator),
        )
    }

    fn validator(&self, id: PeerId) -> Option<&dyn Validator> {
        self.validators.get(&id).map(|v| v as &dyn Validator)
    }

    fn validator_mut(&mut self, id: PeerId) -> Option<&mut dyn Validator> {
        self.validators
            .get_mut(&id)
            .map(|v| v as &mut dyn Validator)
    }

    fn upgrade_validator(&mut self, id: PeerId, version: &Version) -> Result<()> {
        let version = self
            .versions
            .get(version)
            .cloned()
            .ok_or_else(|| anyhow!("Invalid version: {:?}", version))?;
        let validator = self
            .validators
            .get_mut(&id)
            .ok_or_else(|| anyhow!("Invalid id: {}", id))?;
        validator.upgrade(version)
    }

    fn full_nodes<'a>(&'a self) -> Box<dyn Iterator<Item = &'a dyn FullNode> + 'a> {
        Box::new(self.fullnodes.values().map(|v| v as &'a dyn FullNode))
    }

    fn full_nodes_mut<'a>(&'a mut self) -> Box<dyn Iterator<Item = &'a mut dyn FullNode> + 'a> {
        Box::new(
            self.fullnodes
                .values_mut()
                .map(|v| v as &'a mut dyn FullNode),
        )
    }

    fn full_node(&self, id: PeerId) -> Option<&dyn FullNode> {
        self.fullnodes.get(&id).map(|v| v as &dyn FullNode)
    }

    fn full_node_mut(&mut self, id: PeerId) -> Option<&mut dyn FullNode> {
        self.fullnodes.get_mut(&id).map(|v| v as &mut dyn FullNode)
    }

    fn add_validator(&mut self, _version: &Version, _template: NodeConfig) -> Result<PeerId> {
        todo!()
    }

    fn remove_validator(&mut self, _id: PeerId) -> Result<()> {
        todo!()
    }

    fn add_full_node(&mut self, version: &Version, template: NodeConfig) -> Result<PeerId> {
        self.add_fullnode(version, template)
    }

    fn remove_full_node(&mut self, id: PeerId) -> Result<()> {
        if let Some(mut fullnode) = self.fullnodes.remove(&id) {
            fullnode.stop();
        }

        Ok(())
    }

    fn versions<'a>(&'a self) -> Box<dyn Iterator<Item = Version> + 'a> {
        Box::new(self.versions.keys().cloned())
    }

    fn chain_info(&mut self) -> ChainInfo<'_> {
        ChainInfo::new(
            &mut self.root_account,
            &mut self.treasury_compliance_account,
            &mut self.designated_dealer_account,
            self.validators
                .values()
                .map(|v| v.json_rpc_endpoint().to_string())
                .next()
                .unwrap(),
            self.chain_id,
        )
    }

    fn logs_location(&mut self) -> String {
        self.dir.persist();
        self.dir.display().to_string()
    }
}
<commit_msg>forge: add getters for validators for LocalSwarm<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    ChainInfo, FullNode, HealthCheckError, LocalNode, LocalVersion, Node, Swarm, SwarmExt,
    Validator, Version,
};
use anyhow::{anyhow, Result};
use diem_config::config::NodeConfig;
use diem_genesis_tool::{fullnode_builder::FullnodeConfig, validator_builder::ValidatorBuilder};
use diem_sdk::{
    crypto::ed25519::Ed25519PrivateKey,
    types::{
        chain_id::ChainId, transaction::Transaction, waypoint::Waypoint, AccountKey, LocalAccount,
        PeerId,
    },
};
use std::{
    collections::HashMap,
    convert::TryFrom,
    fs, mem,
    num::NonZeroUsize,
    ops,
    path::{Path, PathBuf},
    sync::Arc,
    time::{Duration, Instant},
};
use tempfile::TempDir;

#[derive(Debug)]
pub enum SwarmDirectory {
    Persistent(PathBuf),
    Temporary(TempDir),
}

impl SwarmDirectory {
    pub fn persist(&mut self) {
        match self {
            SwarmDirectory::Persistent(_) => {}
            SwarmDirectory::Temporary(_) => {
                let mut temp = SwarmDirectory::Persistent(PathBuf::new());
                mem::swap(self, &mut temp);
                let _ = mem::replace(self, temp.into_persistent());
            }
        }
    }

    pub fn into_persistent(self) -> Self {
        match self {
            SwarmDirectory::Temporary(tempdir) => SwarmDirectory::Persistent(tempdir.into_path()),
            SwarmDirectory::Persistent(dir) => SwarmDirectory::Persistent(dir),
        }
    }
}

impl ops::Deref for SwarmDirectory {
    type Target = Path;

    fn deref(&self) -> &Self::Target {
        match self {
            SwarmDirectory::Persistent(dir) => dir.deref(),
            SwarmDirectory::Temporary(dir) => dir.path(),
        }
    }
}

impl AsRef<Path> for SwarmDirectory {
    fn as_ref(&self) -> &Path {
        match self {
            SwarmDirectory::Persistent(dir) => dir.as_ref(),
            SwarmDirectory::Temporary(dir) => dir.as_ref(),
        }
    }
}

#[derive(Debug)]
struct ValidatorNetworkAddressEncryptionKey {
    key: diem_sdk::types::network_address::encrypted::Key,
    version: diem_sdk::types::network_address::encrypted::KeyVersion,
}

pub struct LocalSwarmBuilder {
    versions: Arc<HashMap<Version, LocalVersion>>,
    initial_version: Option<Version>,
    template: NodeConfig,
    number_of_validators: NonZeroUsize,
    dir: Option<PathBuf>,
}

impl LocalSwarmBuilder {
    pub fn new(versions: Arc<HashMap<Version, LocalVersion>>) -> Self {
        Self {
            versions,
            initial_version: None,
            template: NodeConfig::default_for_validator(),
            number_of_validators: NonZeroUsize::new(1).unwrap(),
            dir: None,
        }
    }

    pub fn initial_version(mut self, initial_version: Version) -> Self {
        self.initial_version = Some(initial_version);
        self
    }

    pub fn template(mut self, template: NodeConfig) -> Self {
        self.template = template;
        self
    }

    pub fn number_of_validators(mut self, number_of_validators: NonZeroUsize) -> Self {
        self.number_of_validators = number_of_validators;
        self
    }

    pub fn dir<T: AsRef<Path>>(mut self, dir: T) -> Self {
        self.dir = Some(dir.as_ref().into());
        self
    }

    pub fn build<R>(self, rng: R) -> Result<LocalSwarm>
    where
        R: ::rand::RngCore + ::rand::CryptoRng,
    {
        let dir = if let Some(dir) = self.dir {
            if dir.exists() {
                fs::remove_dir_all(&dir)?;
            }
            fs::create_dir_all(&dir)?;
            SwarmDirectory::Persistent(dir)
        } else {
            SwarmDirectory::Temporary(TempDir::new()?)
        };

        let (root_keys, genesis, genesis_waypoint, validators) = ValidatorBuilder::new(
            &dir,
            diem_framework_releases::current_module_blobs().to_vec(),
        )
        .num_validators(self.number_of_validators)
        .template(self.template)
        .build(rng)?;

        // Get the initial version to start the nodes with, either the one provided or fallback to
        // using the the latest version
        let versions = self.versions;
        let initial_version = self.initial_version.unwrap_or_else(|| {
            versions
                .iter()
                .max_by(|v1, v2| v1.0.cmp(v2.0))
                .unwrap()
                .0
                .clone()
        });
        let version = versions.get(&initial_version).unwrap();

        let validators = validators
            .into_iter()
            .map(|v| {
                let node = LocalNode::new(version.to_owned(), v.name, v.directory)?;
                Ok((node.peer_id(), node))
            })
            .collect::<Result<HashMap<_, _>>>()?;

        let validator_network_address_encryption_key = ValidatorNetworkAddressEncryptionKey {
            key: root_keys.validator_network_address_encryption_key,
            version: root_keys.validator_network_address_encryption_key_version,
        };

        let root_account = LocalAccount::new(
            diem_sdk::types::account_config::diem_root_address(),
            AccountKey::from_private_key(root_keys.root_key),
            1,
        );

        // Designated dealer account is built to have the same private key as the TC account
        let designated_dealer_account = LocalAccount::new(
            diem_sdk::types::account_config::testnet_dd_account_address(),
            AccountKey::from_private_key(
                Ed25519PrivateKey::try_from(root_keys.treasury_compliance_key.to_bytes().as_ref())
                    .unwrap(),
            ),
            0,
        );

        let treasury_compliance_account = LocalAccount::new(
            diem_sdk::types::account_config::treasury_compliance_account_address(),
            AccountKey::from_private_key(root_keys.treasury_compliance_key),
            0,
        );

        Ok(LocalSwarm {
            node_name_counter: validators.len() as u64,
            genesis,
            genesis_waypoint,
            versions,
            validators,
            fullnodes: HashMap::new(),
            dir,
            validator_network_address_encryption_key,
            root_account,
            treasury_compliance_account,
            designated_dealer_account,
            chain_id: ChainId::test(),
        })
    }
}

#[derive(Debug)]
pub struct LocalSwarm {
    node_name_counter: u64,
    genesis: Transaction,
    genesis_waypoint: Waypoint,
    versions: Arc<HashMap<Version, LocalVersion>>,
    validators: HashMap<PeerId, LocalNode>,
    fullnodes: HashMap<PeerId, LocalNode>,
    dir: SwarmDirectory,
    validator_network_address_encryption_key: ValidatorNetworkAddressEncryptionKey,
    root_account: LocalAccount,
    treasury_compliance_account: LocalAccount,
    designated_dealer_account: LocalAccount,
    chain_id: ChainId,
}

impl LocalSwarm {
    pub fn builder(versions: Arc<HashMap<Version, LocalVersion>>) -> LocalSwarmBuilder {
        LocalSwarmBuilder::new(versions)
    }

    pub fn launch(&mut self) -> Result<()> {
        // Start all the validators
        for validator in self.validators.values_mut() {
            validator.start()?;
        }

        // Wait for all of them to startup
        let deadline = Instant::now() + Duration::from_secs(60);
        self.wait_for_startup()?;
        self.wait_for_connectivity(deadline)?;
        self.liveness_check(deadline)?;

        Ok(())
    }

    fn wait_for_startup(&mut self) -> Result<()> {
        let num_attempts = 10;
        let mut done = vec![false; self.validators.len()];
        for i in 0..num_attempts {
            println!("Wait for startup attempt: {} of {}", i, num_attempts);
            for (node, done) in self.validators.values_mut().zip(done.iter_mut()) {
                if *done {
                    continue;
                }
                match node.health_check() {
                    Ok(()) => *done = true,

                    Err(HealthCheckError::Unknown(e)) => {
                        return Err(anyhow!("Diem node '{}' is not running: {}", node.name(), e));
                    }
                    Err(HealthCheckError::NotRunning) => {
                        return Err(anyhow!("Diem node '{}' is not running", node.name()));
                    }
                    Err(HealthCheckError::RpcFailure(e)) => {
                        println!("rpc error: {}", e);
                        continue;
                    }
                }
            }

            // Check if all the nodes have been successfully launched
            if done.iter().all(|status| *status) {
                return Ok(());
            }

            ::std::thread::sleep(::std::time::Duration::from_millis(1000));
        }

        Err(anyhow!("Launching Swarm timed out"))
    }

    fn add_fullnode(&mut self, version: &Version, template: NodeConfig) -> Result<PeerId> {
        let name = self.node_name_counter.to_string();
        self.node_name_counter += 1;
        let fullnode_config = FullnodeConfig::public_fullnode(
            name,
            self.dir.as_ref(),
            template,
            &self.genesis_waypoint,
            &self.genesis,
        )?;

        let version = self.versions.get(version).unwrap();
        let mut fullnode = LocalNode::new(
            version.to_owned(),
            fullnode_config.name,
            fullnode_config.directory,
        )?;

        let peer_id = fullnode.peer_id();
        fullnode.start()?;

        self.fullnodes.insert(peer_id, fullnode);

        Ok(peer_id)
    }

    pub fn chain_id(&self) -> ChainId {
        self.chain_id
    }

    pub fn validator(&self, peer_id: PeerId) -> Option<&LocalNode> {
        self.validators.get(&peer_id)
    }

    pub fn validator_mut(&mut self, peer_id: PeerId) -> Option<&mut LocalNode> {
        self.validators.get_mut(&peer_id)
    }

    pub fn validators(&self) -> impl Iterator<Item = &LocalNode> {
        self.validators.values()
    }

    pub fn validators_mut(&mut self) -> impl Iterator<Item = &mut LocalNode> {
        self.validators.values_mut()
    }

    pub fn dir(&self) -> &Path {
        self.dir.as_ref()
    }
}

impl Swarm for LocalSwarm {
    fn health_check(&mut self) -> Result<()> {
        todo!()
    }

    fn validators<'a>(&'a self) -> Box<dyn Iterator<Item = &'a dyn Validator> + 'a> {
        Box::new(self.validators.values().map(|v| v as &'a dyn Validator))
    }

    fn validators_mut<'a>(&'a mut self) -> Box<dyn Iterator<Item = &'a mut dyn Validator> + 'a> {
        Box::new(
            self.validators
                .values_mut()
                .map(|v| v as &'a mut dyn Validator),
        )
    }

    fn validator(&self, id: PeerId) -> Option<&dyn Validator> {
        self.validators.get(&id).map(|v| v as &dyn Validator)
    }

    fn validator_mut(&mut self, id: PeerId) -> Option<&mut dyn Validator> {
        self.validators
            .get_mut(&id)
            .map(|v| v as &mut dyn Validator)
    }

    fn upgrade_validator(&mut self, id: PeerId, version: &Version) -> Result<()> {
        let version = self
            .versions
            .get(version)
            .cloned()
            .ok_or_else(|| anyhow!("Invalid version: {:?}", version))?;
        let validator = self
            .validators
            .get_mut(&id)
            .ok_or_else(|| anyhow!("Invalid id: {}", id))?;
        validator.upgrade(version)
    }

    fn full_nodes<'a>(&'a self) -> Box<dyn Iterator<Item = &'a dyn FullNode> + 'a> {
        Box::new(self.fullnodes.values().map(|v| v as &'a dyn FullNode))
    }

    fn full_nodes_mut<'a>(&'a mut self) -> Box<dyn Iterator<Item = &'a mut dyn FullNode> + 'a> {
        Box::new(
            self.fullnodes
                .values_mut()
                .map(|v| v as &'a mut dyn FullNode),
        )
    }

    fn full_node(&self, id: PeerId) -> Option<&dyn FullNode> {
        self.fullnodes.get(&id).map(|v| v as &dyn FullNode)
    }

    fn full_node_mut(&mut self, id: PeerId) -> Option<&mut dyn FullNode> {
        self.fullnodes.get_mut(&id).map(|v| v as &mut dyn FullNode)
    }

    fn add_validator(&mut self, _version: &Version, _template: NodeConfig) -> Result<PeerId> {
        todo!()
    }

    fn remove_validator(&mut self, _id: PeerId) -> Result<()> {
        todo!()
    }

    fn add_full_node(&mut self, version: &Version, template: NodeConfig) -> Result<PeerId> {
        self.add_fullnode(version, template)
    }

    fn remove_full_node(&mut self, id: PeerId) -> Result<()> {
        if let Some(mut fullnode) = self.fullnodes.remove(&id) {
            fullnode.stop();
        }

        Ok(())
    }

    fn versions<'a>(&'a self) -> Box<dyn Iterator<Item = Version> + 'a> {
        Box::new(self.versions.keys().cloned())
    }

    fn chain_info(&mut self) -> ChainInfo<'_> {
        ChainInfo::new(
            &mut self.root_account,
            &mut self.treasury_compliance_account,
            &mut self.designated_dealer_account,
            self.validators
                .values()
                .map(|v| v.json_rpc_endpoint().to_string())
                .next()
                .unwrap(),
            self.chain_id,
        )
    }

    fn logs_location(&mut self) -> String {
        self.dir.persist();
        self.dir.display().to_string()
    }
}
<|endoftext|>


<commit_before>use std::path::Path;
use std::sync::{Arc, Mutex};
use std::{fs::create_dir_all, path::PathBuf};

use anyhow::{bail, Result};
use crossbeam_channel::unbounded;
use log::warn;

use pueue_lib::network::certificate::create_certificates;
use pueue_lib::network::message::Message;
use pueue_lib::network::protocol::socket_cleanup;
use pueue_lib::network::secret::init_shared_secret;
use pueue_lib::settings::Settings;
use pueue_lib::state::State;

use crate::network::socket::accept_incoming;
use crate::task_handler::TaskHandler;

pub mod cli;
mod network;
mod pid;
mod platform;
mod task_handler;

/// The main entry point for the daemon logic.
/// It's basically the `main`, but publicly exported as a library.
/// That way we can properly do integration testing for the daemon.
///
/// For the purpose of testing, some things shouldn't be run during tests.
/// There are some global operations that crash during tests, such as the ctlc handler.
/// This is due to the fact, that tests in the same file are executed in multiple threads.
/// Since the threads own the same global space, this would crash.
pub async fn run(config_path: Option<PathBuf>, test: bool) -> Result<()> {
    // Try to read settings from the configuration file.
    let settings = match Settings::read(&config_path) {
        Ok(settings) => settings,
        Err(_) => {
            // There's something wrong with the config file or something's missing.
            // Try to read the config and fill missing values with defaults.
            // This might be possible on version upgrade or first run.
            let settings = Settings::read_with_defaults(false, &config_path)?;

            // Since we needed to add values to the configuration, we have to save it.
            // This also creates the save file in case it didn't exist yet.
            if let Err(error) = settings.save(&config_path) {
                bail!("Failed saving config file: {:?}.", error);
            }
            settings
        }
    };

    init_directories(&settings.shared.pueue_directory());
    if !settings.shared.daemon_key().exists() && !settings.shared.daemon_cert().exists() {
        create_certificates(&settings.shared)?;
    }
    init_shared_secret(&settings.shared.shared_secret_path())?;
    pid::create_pid_file(&settings.shared.pueue_directory())?;

    let mut state = State::new(&settings, config_path.clone());
    // Restore the previous state and save any changes that might have happened during this process
    if let Err(error) = state.restore() {
        warn!("Failed to restore previous state:\n {:?}", error);
        warn!("Using clean state instead.");
    };
    state.save()?;
    let state = Arc::new(Mutex::new(state));

    let (sender, receiver) = unbounded();
    let mut task_handler = TaskHandler::new(state.clone(), receiver);

    // Don't set ctrlc handler during testing.
    // This is necessary for multithreaded integration testing.
    // On top of this, ctrlc also somehow breaks test error output.
    if !test {
        // This section handles Shutdown via SigTerm/SigInt process signals
        // 1. Remove the unix socket (if it exists).
        // 2. Notify the TaskHandler, so it can shutdown gracefully.
        //
        // The actual program exit will be done via the TaskHandler.
        let sender_clone = sender.clone();
        let settings_clone = settings.clone();
        ctrlc::set_handler(move || {
            if let Err(error) = socket_cleanup(&settings_clone.shared) {
                println!("Failed to cleanup socket after shutdown signal.");
                println!("{}", error);
            }

            // Notify the task handler
            sender_clone
                .send(Message::DaemonShutdown)
                .expect("Failed to send Message to TaskHandler on Shutdown");
        })?;
    }

    let orig_hook = std::panic::take_hook();
    let settings_clone = settings.clone();

    // Don't set panic hook during testing.
    // This is necessary for multithreaded integration testing.
    if !test {
        std::panic::set_hook(Box::new(move |panic_info| {
            // invoke the default handler and exit the process
            orig_hook(panic_info);
            if let Err(error) = pid::cleanup_pid_file(&settings_clone.shared.pueue_directory()) {
                println!("Failed to cleanup pid after panic.");
                println!("{}", error);
            }
            if let Err(error) = socket_cleanup(&settings_clone.shared) {
                println!("Failed to cleanup socket after panic.");
                println!("{}", error);
            }

            std::process::exit(1);
        }));
    }

    std::thread::spawn(move || {
        task_handler.run();
    });

    accept_incoming(sender, state.clone()).await?;

    Ok(())
}

/// Initialize all directories needed for normal operation.
fn init_directories(pueue_dir: &Path) {
    // Pueue base path
    if !pueue_dir.exists() {
        if let Err(error) = create_dir_all(&pueue_dir) {
            panic!(
                "Failed to create main directory at {:?} error: {:?}",
                pueue_dir, error
            );
        }
    }

    // Task log dir
    let log_dir = pueue_dir.join("log");
    if !log_dir.exists() {
        if let Err(error) = create_dir_all(&log_dir) {
            panic!(
                "Failed to create log directory at {:?} error: {:?}",
                log_dir, error
            );
        }
    }

    // Task certs dir
    let certs_dir = pueue_dir.join("certs");
    if !certs_dir.exists() {
        if let Err(error) = create_dir_all(&certs_dir) {
            panic!(
                "Failed to create certificate directory at {:?} error: {:?}",
                certs_dir, error
            );
        }
    }

    // Task log dir
    let logs_dir = pueue_dir.join("task_logs");
    if !logs_dir.exists() {
        if let Err(error) = create_dir_all(&logs_dir) {
            panic!(
                "Failed to create task logs directory at {:?} error: {:?}",
                logs_dir, error
            );
        }
    }
}
<commit_msg>Move signal/panic initialization to its own function<commit_after>use std::path::Path;
use std::sync::{Arc, Mutex};
use std::{fs::create_dir_all, path::PathBuf};

use anyhow::{bail, Result};
use crossbeam_channel::{unbounded, Sender};
use log::warn;

use pueue_lib::network::certificate::create_certificates;
use pueue_lib::network::message::Message;
use pueue_lib::network::protocol::socket_cleanup;
use pueue_lib::network::secret::init_shared_secret;
use pueue_lib::settings::Settings;
use pueue_lib::state::State;

use crate::network::socket::accept_incoming;
use crate::task_handler::TaskHandler;

pub mod cli;
mod network;
mod pid;
mod platform;
mod task_handler;

/// The main entry point for the daemon logic.
/// It's basically the `main`, but publicly exported as a library.
/// That way we can properly do integration testing for the daemon.
///
/// For the purpose of testing, some things shouldn't be run during tests.
/// There are some global operations that crash during tests, such as the ctlc handler.
/// This is due to the fact, that tests in the same file are executed in multiple threads.
/// Since the threads own the same global space, this would crash.
pub async fn run(config_path: Option<PathBuf>, test: bool) -> Result<()> {
    // Try to read settings from the configuration file.
    let settings = match Settings::read(&config_path) {
        Ok(settings) => settings,
        Err(_) => {
            // There's something wrong with the config file or something's missing.
            // Try to read the config and fill missing values with defaults.
            // This might be possible on version upgrade or first run.
            let settings = Settings::read_with_defaults(false, &config_path)?;

            // Since we needed to add values to the configuration, we have to save it.
            // This also creates the save file in case it didn't exist yet.
            if let Err(error) = settings.save(&config_path) {
                bail!("Failed saving config file: {:?}.", error);
            }
            settings
        }
    };

    init_directories(&settings.shared.pueue_directory());
    if !settings.shared.daemon_key().exists() && !settings.shared.daemon_cert().exists() {
        create_certificates(&settings.shared)?;
    }
    init_shared_secret(&settings.shared.shared_secret_path())?;
    pid::create_pid_file(&settings.shared.pueue_directory())?;

    let mut state = State::new(&settings, config_path.clone());
    // Restore the previous state and save any changes that might have happened during this process
    if let Err(error) = state.restore() {
        warn!("Failed to restore previous state:\n {:?}", error);
        warn!("Using clean state instead.");
    };
    state.save()?;
    let state = Arc::new(Mutex::new(state));

    let (sender, receiver) = unbounded();
    let mut task_handler = TaskHandler::new(state.clone(), receiver);

    // Don't set ctrlc and panic handlers during testing.
    // This is necessary for multithreaded integration testing, since multiple listener per process
    // aren't prophibited. On top of this, ctrlc also somehow breaks test error output.
    if !test {
        setup_signal_panic_handling(&settings, &sender)?;
    }

    std::thread::spawn(move || {
        task_handler.run();
    });

    accept_incoming(sender, state.clone()).await?;

    Ok(())
}

/// Initialize all directories needed for normal operation.
fn init_directories(pueue_dir: &Path) {
    // Pueue base path
    if !pueue_dir.exists() {
        if let Err(error) = create_dir_all(&pueue_dir) {
            panic!(
                "Failed to create main directory at {:?} error: {:?}",
                pueue_dir, error
            );
        }
    }

    // Task log dir
    let log_dir = pueue_dir.join("log");
    if !log_dir.exists() {
        if let Err(error) = create_dir_all(&log_dir) {
            panic!(
                "Failed to create log directory at {:?} error: {:?}",
                log_dir, error
            );
        }
    }

    // Task certs dir
    let certs_dir = pueue_dir.join("certs");
    if !certs_dir.exists() {
        if let Err(error) = create_dir_all(&certs_dir) {
            panic!(
                "Failed to create certificate directory at {:?} error: {:?}",
                certs_dir, error
            );
        }
    }

    // Task log dir
    let logs_dir = pueue_dir.join("task_logs");
    if !logs_dir.exists() {
        if let Err(error) = create_dir_all(&logs_dir) {
            panic!(
                "Failed to create task logs directory at {:?} error: {:?}",
                logs_dir, error
            );
        }
    }
}

/// Setup signal handling and panic handling.
///
/// On SIGINT and SIGTERM, we exit gracefully by sending a DaemonShutdown message to the
/// TaskHandler. This is to prevent dangling processes and other weird edge-cases.
///
/// On panic, we want to cleanup existing unix sockets and the PID file.
fn setup_signal_panic_handling(settings: &Settings, sender: &Sender<Message>) -> Result<()> {
    let sender_clone = sender.clone();
    let settings_clone = settings.clone();

    // This section handles Shutdown via SigTerm/SigInt process signals
    // 1. Remove the unix socket (if it exists).
    // 2. Notify the TaskHandler, so it can shutdown gracefully.
    //
    // The actual program exit will be done via the TaskHandler.
    ctrlc::set_handler(move || {
        if let Err(error) = socket_cleanup(&settings_clone.shared) {
            println!("Failed to cleanup socket after shutdown signal.");
            println!("{}", error);
        }

        // Notify the task handler
        sender_clone
            .send(Message::DaemonShutdown)
            .expect("Failed to send Message to TaskHandler on Shutdown");
    })?;

    let settings_clone = settings.clone();
    let orig_hook = std::panic::take_hook();
    std::panic::set_hook(Box::new(move |panic_info| {
        // invoke the default handler and exit the process
        orig_hook(panic_info);
        if let Err(error) = pid::cleanup_pid_file(&settings_clone.shared.pueue_directory()) {
            println!("Failed to cleanup pid after panic.");
            println!("{}", error);
        }
        if let Err(error) = socket_cleanup(&settings_clone.shared) {
            println!("Failed to cleanup socket after panic.");
            println!("{}", error);
        }

        std::process::exit(1);
    }));

    Ok(())
}
<|endoftext|>


<commit_before>/*
--- Day 4: Security Through Obscurity ---

Finally, you come across an information kiosk with a list of rooms. Of course, the list is encrypted and full of decoy data, but the instructions to decode the list are barely hidden nearby. Better remove the decoy data first.

Each room consists of an encrypted name (lowercase letters separated by dashes) followed by a dash, a sector ID, and a checksum in square brackets.

A room is real (not a decoy) if the checksum is the five most common letters in the encrypted name, in order, with ties broken by alphabetization. For example:

aaaaa-bbb-z-y-x-123[abxyz] is a real room because the most common letters are a (5), b (3), and then a tie between x, y, and z, which are listed alphabetically.
a-b-c-d-e-f-g-h-987[abcde] is a real room because although the letters are all tied (1 of each), the first five are listed alphabetically.
not-a-real-room-404[oarel] is a real room.
totally-real-room-200[decoy] is not.
Of the real rooms from the list above, the sum of their sector IDs is 1514.

What is the sum of the sector IDs of the real rooms?

*/

extern crate regex;
#[macro_use] extern crate lazy_static;
use regex::Regex;
use std::io::prelude::*;
use std::fs::File;

#[derive(Debug)]
struct Room {
    name : String,
    sector_id : isize,
    checksum: String,
}

impl Room {
    fn from_line(line : &str) -> Room {

        lazy_static! {
            // compile regex only once
            static ref REGEX : Regex = Regex::new(r"^((([a-zA-Z0-9]+\-)+)([0-9]+)\[(.+)\])$").unwrap();
        }

        // find matches in line
        let cap = REGEX.captures(line).unwrap();

        // grab the parts we need
        let name : String     = cap.get(2).unwrap().as_str().to_string();
        let sector_id : isize = cap.get(4).unwrap().as_str().parse().unwrap();
        let checksum : String = cap.get(5).unwrap().as_str().to_string();

        // and create a Room struct off of it
        Room { name: name, sector_id: sector_id, checksum: checksum }
    }

    // does checksum match input?
    fn is_valid(&self) -> bool {
        // remove the dashes
        let undashed_name : &String = &self.name.replace('-', "");

        // grab the chars as an array
        let mut chars_as_vec : Vec<char> = undashed_name.chars().collect();

        // sort the array alphabetically
        chars_as_vec.sort();

        // dedupe the array
        chars_as_vec.dedup();

        // create tuples with the character and numer of occurances in original string
        let mut char_and_count : Vec<(char, usize)> = chars_as_vec.iter()
            .map(|c| { (*c, undashed_name.matches(*c).count()) })
            .collect();

        // ordering.. magic :)
        char_and_count.sort_by(|a, b| {
            // if a.1 (number of occurances) equal..
            if a.1 == b.1 {
                // matching; compare using the char alphabetically
                return b.0.cmp(&a.0)
            }
            // not matching; how do they compare?
            a.1.cmp(&b.1)
        });

        // since the order is backwards, lets reverse
        char_and_count.reverse();

        // truncate to first five chars
        char_and_count.truncate(5);

        // grab the chars and concatenate as string
        let checksum_as_vec : Vec<char> = char_and_count.iter().map(|c| c.0).collect();
        let checksum : String = checksum_as_vec.clone().into_iter().collect();

        // a room is valid once the two checksums match
        return checksum == self.checksum
    }
}

fn main() {
    let mut f = File::open("input.txt").unwrap();
    let mut input = String::new();
    f.read_to_string(&mut input).ok();

    let rooms : Vec<Room> = input.lines().map(|line| {Room::from_line(line)}).collect();

    // assert our test sequence
    let sum = rooms.iter().filter(|r| {r.is_valid()}).fold(0, |sum, r| { sum + r.sector_id});
    println!("sum: {}", sum);
}
<commit_msg>Copy changes<commit_after>/*
--- Day 4: Security Through Obscurity ---

Finally, you come across an information kiosk with a list of rooms. Of course, the list is encrypted and full of decoy data, but the instructions to decode the list are barely hidden nearby. Better remove the decoy data first.

Each room consists of an encrypted name (lowercase letters separated by dashes) followed by a dash, a sector ID, and a checksum in square brackets.

A room is real (not a decoy) if the checksum is the five most common letters in the encrypted name, in order, with ties broken by alphabetization. For example:

aaaaa-bbb-z-y-x-123[abxyz] is a real room because the most common letters are a (5), b (3), and then a tie between x, y, and z, which are listed alphabetically.
a-b-c-d-e-f-g-h-987[abcde] is a real room because although the letters are all tied (1 of each), the first five are listed alphabetically.
not-a-real-room-404[oarel] is a real room.
totally-real-room-200[decoy] is not.
Of the real rooms from the list above, the sum of their sector IDs is 1514.

What is the sum of the sector IDs of the real rooms?

*/

extern crate regex;
#[macro_use] extern crate lazy_static;
use regex::Regex;
use std::io::prelude::*;
use std::fs::File;

#[derive(Debug)]
struct Room {
    name : String,
    sector_id : isize,
    checksum: String,
}

impl Room {
    fn from_line(line : &str) -> Room {

        lazy_static! {
            // compile regex only once
            static ref REGEX : Regex = Regex::new(r"^((([a-zA-Z0-9]+\-)+)([0-9]+)\[(.+)\])$").unwrap();
        }

        // find matches in line
        let cap = REGEX.captures(line).unwrap();

        // grab the parts we need
        let name : String     = cap.get(2).unwrap().as_str().to_string();
        let sector_id : isize = cap.get(4).unwrap().as_str().parse().unwrap();
        let checksum : String = cap.get(5).unwrap().as_str().to_string();

        // and create a Room struct off of it
        Room { name: name, sector_id: sector_id, checksum: checksum }
    }

    // does checksum match input?
    fn is_valid(&self) -> bool {
        // remove the dashes
        let undashed_name : &String = &self.name.replace('-', "");

        // grab the chars as an array
        let mut chars_as_vec : Vec<char> = undashed_name.chars().collect();

        // sort the array alphabetically
        chars_as_vec.sort();

        // dedupe the array
        chars_as_vec.dedup();

        // create tuples with the character and numer of occurances in original string
        let mut char_and_count : Vec<(char, usize)> = chars_as_vec.iter()
            .map(|c| { (*c, undashed_name.matches(*c).count()) })
            .collect();

        // ordering.. magic :)
        char_and_count.sort_by(|a, b| {
            // if a.1 (number of occurances) equal..
            if a.1 == b.1 {
                // matching; compare using the char alphabetically
                return b.0.cmp(&a.0)
            }
            // not matching; how do they compare?
            a.1.cmp(&b.1)
        });

        // since the order is backwards, lets reverse
        char_and_count.reverse();

        // truncate to first five chars
        char_and_count.truncate(5);

        // grab the chars and concatenate as string
        let checksum_as_vec : Vec<char> = char_and_count.iter().map(|c| c.0).collect();
        let checksum : String = checksum_as_vec.clone().into_iter().collect();

        // a room is valid once the two checksums match
        return checksum == self.checksum
    }
}

fn main() {
    let mut f = File::open("input.txt").unwrap();
    let mut input = String::new();
    f.read_to_string(&mut input).ok();

    let rooms : Vec<Room> = input.lines().map(|line| {Room::from_line(line)}).collect();

    // grab the valid rooms, and sum the sector id using fold
    let sum = rooms.iter()
        .filter(|r| {r.is_valid()})
        .fold(0, |sum, r| { sum + r.sector_id});

    // print the answer 🎉
    println!("sum: {}", sum);
}
<|endoftext|>


<commit_before>use cargo::core::{enable_nightly_features, Package, SourceId, Workspace};
use cargo::sources::SourceConfigMap;
use cargo::util::{internal, Config};
use error::Result;
use rustwide::{cmd::SandboxBuilder, Crate, Toolchain, WorkspaceBuilder};
use std::collections::HashSet;
use std::path::Path;
use utils::{get_current_versions, parse_rustc_version, resolve_deps};
use Metadata;

pub fn build_doc_rustwide(name: &str, version: &str, target: Option<&str>) -> Result<Package> {
    // TODO: Handle workspace path correctly
    let rustwide_workspace =
        WorkspaceBuilder::new(Path::new("/tmp/docs-builder"), "docsrs").init()?;

    // TODO: Instead of using just nightly, we can pin a version.
    //       Docs.rs can only use nightly (due to unstable docs.rs features in rustdoc)
    let toolchain = Toolchain::Dist {
        name: "nightly".into(),
    };
    toolchain.install(&rustwide_workspace)?;

    let krate = Crate::crates_io(name, version);
    krate.fetch(&rustwide_workspace)?;

    // Configure a sandbox with 1GB of RAM and no network access
    // TODO: 1GB might not be enough
    let sandbox = SandboxBuilder::new()
        .memory_limit(Some(1024 * 1024 * 1024))
        .enable_networking(false);

    let mut build_dir = rustwide_workspace.build_dir(&format!("{}-{}", name, version));
    let pkg = build_dir.build(&toolchain, &krate, sandbox, |build| {
        enable_nightly_features();
        let config = Config::default()?;
        let source_id = try!(SourceId::crates_io(&config));
        let source_cfg_map = try!(SourceConfigMap::new(&config));
        let manifest_path = build.host_source_dir().join("Cargo.toml");
        let ws = Workspace::new(&manifest_path, &config)?;
        let pkg = ws.load(&manifest_path)?;

        let metadata = Metadata::from_package(&pkg).map_err(|e| internal(e.to_string()))?;

        let mut rustdoc_flags: Vec<String> = vec![
            "-Z".to_string(),
            "unstable-options".to_string(),
            "--resource-suffix".to_string(),
            // FIXME: We need to get rustc version inside of container.
            //        Our get_current_versions gets rustc version from host system not container.
            format!("-{}", parse_rustc_version(get_current_versions()?.0)?),
            "--static-root-path".to_string(),
            "/".to_string(),
            "--disable-per-crate-search".to_string(),
        ];

        let source = try!(source_cfg_map.load(source_id, &HashSet::new()));
        let _lock = try!(config.acquire_package_cache_lock());

        for (name, dep) in try!(resolve_deps(&pkg, &config, source)) {
            rustdoc_flags.push("--extern-html-root-url".to_string());
            rustdoc_flags.push(format!(
                "{}=https://docs.rs/{}/{}",
                name.replace("-", "_"),
                dep.name(),
                dep.version()
            ));
        }

        let mut cargo_args = vec!["doc".to_owned(), "--lib".to_owned(), "--no-deps".to_owned()];
        if let Some(features) = &metadata.features {
            cargo_args.push("--features".to_owned());
            cargo_args.push(features.join(" "));
        }
        if metadata.all_features {
            cargo_args.push("--all-features".to_owned());
        }
        if metadata.no_default_features {
            cargo_args.push("--no-default-features".to_owned());
        }

        // TODO: We need to use build result here
        // FIXME: We also need build log (basically stderr message)
        let result = build
            .cargo()
            .env(
                "RUSTFLAGS",
                metadata
                    .rustc_args
                    .map(|args| args.join(""))
                    .unwrap_or("".to_owned()),
            )
            .env("RUSTDOCFLAGS", rustdoc_flags.join(" "))
            .args(&cargo_args)
            .run();

        // TODO: We need to return build result as well
        Ok(pkg)
    })?;

    Ok(pkg)
}
<commit_msg>correctly fetch rustwide's rustc version<commit_after>use cargo::core::{enable_nightly_features, Package, SourceId, Workspace as CargoWorkspace};
use cargo::sources::SourceConfigMap;
use cargo::util::{internal, Config};
use error::Result;
use rustwide::{
    cmd::{Command, SandboxBuilder},
    Crate, Toolchain, Workspace, WorkspaceBuilder,
};
use std::collections::HashSet;
use std::path::Path;
use utils::{parse_rustc_version, resolve_deps};
use Metadata;

pub fn build_doc_rustwide(name: &str, version: &str, target: Option<&str>) -> Result<Package> {
    // TODO: Handle workspace path correctly
    let rustwide_workspace =
        WorkspaceBuilder::new(Path::new("/tmp/docs-builder"), "docsrs").init()?;

    // TODO: Instead of using just nightly, we can pin a version.
    //       Docs.rs can only use nightly (due to unstable docs.rs features in rustdoc)
    let toolchain = Toolchain::Dist {
        name: "nightly".into(),
    };
    toolchain.install(&rustwide_workspace)?;

    let krate = Crate::crates_io(name, version);
    krate.fetch(&rustwide_workspace)?;

    // Configure a sandbox with 1GB of RAM and no network access
    // TODO: 1GB might not be enough
    let sandbox = SandboxBuilder::new()
        .memory_limit(Some(1024 * 1024 * 1024))
        .enable_networking(false);

    let mut build_dir = rustwide_workspace.build_dir(&format!("{}-{}", name, version));
    let pkg = build_dir.build(&toolchain, &krate, sandbox, |build| {
        enable_nightly_features();
        let config = Config::default()?;
        let source_id = try!(SourceId::crates_io(&config));
        let source_cfg_map = try!(SourceConfigMap::new(&config));
        let manifest_path = build.host_source_dir().join("Cargo.toml");
        let ws = CargoWorkspace::new(&manifest_path, &config)?;
        let pkg = ws.load(&manifest_path)?;

        let metadata = Metadata::from_package(&pkg).map_err(|e| internal(e.to_string()))?;

        let mut rustdoc_flags: Vec<String> = vec![
            "-Z".to_string(),
            "unstable-options".to_string(),
            "--resource-suffix".to_string(),
            format!(
                "-{}",
                parse_rustc_version(rustc_version(&rustwide_workspace, &toolchain)?)?
            ),
            "--static-root-path".to_string(),
            "/".to_string(),
            "--disable-per-crate-search".to_string(),
        ];

        let source = try!(source_cfg_map.load(source_id, &HashSet::new()));
        let _lock = try!(config.acquire_package_cache_lock());

        for (name, dep) in try!(resolve_deps(&pkg, &config, source)) {
            rustdoc_flags.push("--extern-html-root-url".to_string());
            rustdoc_flags.push(format!(
                "{}=https://docs.rs/{}/{}",
                name.replace("-", "_"),
                dep.name(),
                dep.version()
            ));
        }

        let mut cargo_args = vec!["doc".to_owned(), "--lib".to_owned(), "--no-deps".to_owned()];
        if let Some(features) = &metadata.features {
            cargo_args.push("--features".to_owned());
            cargo_args.push(features.join(" "));
        }
        if metadata.all_features {
            cargo_args.push("--all-features".to_owned());
        }
        if metadata.no_default_features {
            cargo_args.push("--no-default-features".to_owned());
        }

        // TODO: We need to use build result here
        // FIXME: We also need build log (basically stderr message)
        let result = build
            .cargo()
            .env(
                "RUSTFLAGS",
                metadata
                    .rustc_args
                    .map(|args| args.join(""))
                    .unwrap_or("".to_owned()),
            )
            .env("RUSTDOCFLAGS", rustdoc_flags.join(" "))
            .args(&cargo_args)
            .run();

        // TODO: We need to return build result as well
        Ok(pkg)
    })?;

    Ok(pkg)
}

fn rustc_version(workspace: &Workspace, toolchain: &Toolchain) -> Result<String> {
    let res = Command::new(workspace, toolchain.rustc())
        .args(&["--version"])
        .log_output(false)
        .run_capture()?;

    if let Some(line) = res.stdout_lines().iter().next() {
        Ok(line.clone())
    } else {
        Err(::failure::err_msg(
            "invalid output returned by `rustc --version`",
        ))
    }
}
<|endoftext|>


<commit_before>#![allow(unused_imports, dead_code, unused_variables)]

#[cfg(feature = "pkg-config")]
extern crate pkg_config;
#[cfg(feature = "bindgen")]
extern crate bindgen;

use std::path::PathBuf;
use std::env;
use std::fs;

const SDL2_BUNDLED_VERSION: &str = "2.0.6";

fn main() {
    let target = env::var("TARGET").expect("Cargo build scripts always have TARGET");
    let host = env::var("HOST").expect("Cargo build scripts always have HOST");

    prepare_bindings(&target, &host);

    if get_os_from_triple(&target).unwrap() == "ios" {
        println!("cargo:rustc-flags=-l framework=AVFoundation");
        println!("cargo:rustc-flags=-l framework=AudioToolbox");
        println!("cargo:rustc-flags=-l framework=CoreAudio");
        println!("cargo:rustc-flags=-l framework=CoreGraphics");
        println!("cargo:rustc-flags=-l framework=CoreMotion");
        println!("cargo:rustc-flags=-l framework=Foundation");
        println!("cargo:rustc-flags=-l framework=GameController");
        println!("cargo:rustc-flags=-l framework=OpenGLES");
        println!("cargo:rustc-flags=-l framework=QuartzCore");
        println!("cargo:rustc-flags=-l framework=UIKit");
    }
}

#[cfg(not(feature = "pkg-config"))]
fn build_pkgconfig() -> bool {
    false
}

#[cfg(feature = "pkg-config")]
fn build_pkgconfig() -> bool {
    pkg_config::probe_library("sdl2").is_ok()
}

#[cfg(not(feature = "bindgen"))]
fn prepare_bindings(target: &str, host: &str) {
    add_explicit_linker_flags(get_os_from_triple(target).unwrap());
    let out_path = PathBuf::from(env::var("OUT_DIR").unwrap());
    let crate_path = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap());
    fs::copy(crate_path.join("pregenerated_bindings.rs"), out_path.join("bindings.rs"))
        .expect("Couldn't find pregenerated bindings!");
}

#[cfg(feature = "bindgen")]
fn prepare_bindings(target: &str, host: &str) {
    let mut bindings = bindgen::Builder::default();

    // Set correct target triple when cross-compiling
    if target != host {
        bindings = bindings.clang_arg("-target");
        bindings = bindings.clang_arg(target.clone());
    }

    if let Ok(include_path) = env::var("SDL2_INCLUDE_PATH") {
        bindings = bindings.clang_arg(String::from("-I") + &include_path);
        add_explicit_linker_flags(get_os_from_triple(target).unwrap());
    } else if build_pkgconfig() {
        #[cfg(feature = "pkg-config")]
        for path in &pkg_config::find_library("sdl2").unwrap().include_paths {
            bindings = bindings.clang_arg(String::from("-I") +
                                          &path.clone().into_os_string().into_string().unwrap());
        }
    } else {
        let mut include_path: PathBuf = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap());
        include_path.push(String::from("SDL2-") + SDL2_BUNDLED_VERSION);
        include_path.push("include");
        bindings = bindings.clang_arg(String::from("-I") +
                                      &include_path.into_os_string().into_string().unwrap());
        add_explicit_linker_flags(get_os_from_triple(target).unwrap());
    }

    let bindings = bindings
        .header("wrapper.h")
        .hide_type("_")
        .hide_type("FP_NAN")
        .hide_type("FP_INFINITE")
        .hide_type("FP_ZERO")
        .hide_type("FP_SUBNORMAL")
        .hide_type("FP_NORMAL") // Until https://github.com/rust-lang-nursery/rust-bindgen/issues/687 gets fixed
        .hide_type("max_align_t") // Until https://github.com/rust-lang-nursery/rust-bindgen/issues/550 gets fixed
        .generate()
        .expect("Unable to generate bindings!");

    let out_path = PathBuf::from(env::var("OUT_DIR").unwrap());

    bindings
        .write_to_file(out_path.join("bindings.rs"))
        .expect("Couldn't write bindings!");
}

fn get_os_from_triple(triple: &str) -> Option<&str>
{
    triple.split("-").nth(2)
}

fn add_explicit_linker_flags(target_os: &str) {
    if cfg!(feature = "use_mac_framework") && target_os == "darwin" {
        println!("cargo:rustc-flags=-l framework=SDL2");
    } else {
        println!("cargo:rustc-flags=-l SDL2");
    }
}
<commit_msg>SDL2 has default configurations for all operating systems but Linux<commit_after>#![allow(unused_imports, dead_code, unused_variables)]

#[cfg(feature = "pkg-config")]
extern crate pkg_config;
#[cfg(feature = "bindgen")]
extern crate bindgen;

use std::path::PathBuf;
use std::env;
use std::fs;

const SDL2_BUNDLED_VERSION: &str = "2.0.6";

fn main() {
    let target = env::var("TARGET").expect("Cargo build scripts always have TARGET");
    let host = env::var("HOST").expect("Cargo build scripts always have HOST");

    prepare_bindings(&target, &host);

    if get_os_from_triple(&target).unwrap() == "ios" {
        println!("cargo:rustc-flags=-l framework=AVFoundation");
        println!("cargo:rustc-flags=-l framework=AudioToolbox");
        println!("cargo:rustc-flags=-l framework=CoreAudio");
        println!("cargo:rustc-flags=-l framework=CoreGraphics");
        println!("cargo:rustc-flags=-l framework=CoreMotion");
        println!("cargo:rustc-flags=-l framework=Foundation");
        println!("cargo:rustc-flags=-l framework=GameController");
        println!("cargo:rustc-flags=-l framework=OpenGLES");
        println!("cargo:rustc-flags=-l framework=QuartzCore");
        println!("cargo:rustc-flags=-l framework=UIKit");
    }
}

#[cfg(not(feature = "pkg-config"))]
fn build_pkgconfig() -> bool {
    false
}

#[cfg(feature = "pkg-config")]
fn build_pkgconfig() -> bool {
    pkg_config::probe_library("sdl2").is_ok()
}

#[cfg(not(feature = "bindgen"))]
fn prepare_bindings(target: &str, host: &str) {
    add_explicit_linker_flags(get_os_from_triple(target).unwrap());
    let out_path = PathBuf::from(env::var("OUT_DIR").unwrap());
    let crate_path = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap());
    fs::copy(crate_path.join("pregenerated_bindings.rs"), out_path.join("bindings.rs"))
        .expect("Couldn't find pregenerated bindings!");
}

#[cfg(feature = "bindgen")]
fn prepare_bindings(target: &str, host: &str) {
    let mut bindings = bindgen::Builder::default();

    // Set correct target triple when cross-compiling
    if target != host {
        bindings = bindings.clang_arg("-target");
        bindings = bindings.clang_arg(target.clone());
    }

    if let Ok(include_path) = env::var("SDL2_INCLUDE_PATH") {
        bindings = bindings.clang_arg(String::from("-I") + &include_path);
        add_explicit_linker_flags(get_os_from_triple(target).unwrap());
    } else if build_pkgconfig() {
        #[cfg(feature = "pkg-config")]
        for path in &pkg_config::find_library("sdl2").unwrap().include_paths {
            bindings = bindings.clang_arg(String::from("-I") +
                                          &path.clone().into_os_string().into_string().unwrap());
        }
    } else {
        let mut include_path: PathBuf = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap());
        include_path.push(String::from("SDL2-") + SDL2_BUNDLED_VERSION);
        include_path.push("include");
        bindings = bindings.clang_arg(String::from("-I") +
                                      &include_path.into_os_string().into_string().unwrap());
        // SDL2 hasn't a default configuration for Linux
        if get_os_from_triple(target).unwrap() == "linux" {
            bindings = bindings.clang_arg("-DSDL_VIDEO_DRIVER_X11");
            bindings = bindings.clang_arg("-DSDL_VIDEO_DRIVER_WAYLAND");
        }
        add_explicit_linker_flags(get_os_from_triple(target).unwrap());
    }

    let bindings = bindings
        .header("wrapper.h")
        .hide_type("_")
        .hide_type("FP_NAN")
        .hide_type("FP_INFINITE")
        .hide_type("FP_ZERO")
        .hide_type("FP_SUBNORMAL")
        .hide_type("FP_NORMAL") // Until https://github.com/rust-lang-nursery/rust-bindgen/issues/687 gets fixed
        .hide_type("max_align_t") // Until https://github.com/rust-lang-nursery/rust-bindgen/issues/550 gets fixed
        .generate()
        .expect("Unable to generate bindings!");

    let out_path = PathBuf::from(env::var("OUT_DIR").unwrap());

    bindings
        .write_to_file(out_path.join("bindings.rs"))
        .expect("Couldn't write bindings!");
}

fn get_os_from_triple(triple: &str) -> Option<&str>
{
    triple.split("-").nth(2)
}

fn add_explicit_linker_flags(target_os: &str) {
    if cfg!(feature = "use_mac_framework") && target_os == "darwin" {
        println!("cargo:rustc-flags=-l framework=SDL2");
    } else {
        println!("cargo:rustc-flags=-l SDL2");
    }
}
<|endoftext|>


<commit_before>#![feature(test)]

extern crate intmap;
extern crate rand;
extern crate test;

use intmap::IntMap;
use intmap::IntMap2;
use std::collections::HashMap;
use std::collections::BTreeMap;

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

    const VEC_COUNT: usize = 1000;

    #[bench]
    fn u64_insert_built_in(b: &mut Bencher) {
        let data = get_random_range(VEC_COUNT);
        let mut map = HashMap::with_capacity(data.len());

        b.iter(|| {
            map.clear();

            for s in data.iter() {
               test::black_box(map.insert(s, s));
            }

        });
    }

    #[bench]
    fn u64_get_built_in(b: &mut Bencher) {
        let data = get_random_range(VEC_COUNT);
        let mut map: HashMap<&u64, &u64>  = HashMap::with_capacity(data.len());

        for s in data.iter() {
            test::black_box({
                map.insert(s, s);
            });
        }

        b.iter(|| {
            for s in data.iter() {
                test::black_box({
                    map.contains_key(s);

                });
            }
        });
    }

    #[bench]
    fn u64_insert_intmap(b: &mut Bencher) {
        let data = get_random_range(VEC_COUNT);

        let mut map = IntMap::with_capacity(data.len());

        b.iter(|| {
            map.clear();
            for s in data.iter() {
                test::black_box(map.insert(*s, s));
            }
        });
    }

    #[bench]
    fn u64_get_intmap(b: &mut Bencher) {
        let data = get_random_range(VEC_COUNT);

        let mut map = IntMap::with_capacity(data.len());
        for s in data.iter() {
            map.insert(*s, s);
        }

        b.iter(|| {
            for s in data.iter() {
                test::black_box({
                    map.contains_key(*s);

                });
            }
        });
    }

    #[bench]
    fn u64_insert_intmap2(b: &mut Bencher) {
        let data = get_random_range(VEC_COUNT);

        let mut map = IntMap2::with_capacity(data.len());

        let mut cols = map.collisions();
        let mut vec = Vec::new();

        for c in cols.iter() {
            vec.push(c.clone());

        }
        vec.sort();
        println!("Collisions: {}", vec.len());
        for v in vec.iter() {
            println!("col: {:?}", v);
        }

        b.iter(|| {
            map.clear();
            for s in data.iter() {
                test::black_box(map.insert(*s, s));
            }
        });
    }

    #[bench]
    fn u64_get_intmap2(b: &mut Bencher) {
        let data = get_random_range(VEC_COUNT);

        let mut map = IntMap2::with_capacity(data.len());
        for s in data.iter() {
            map.insert(*s, s);
        }

        b.iter(|| {
            for s in data.iter() {
                test::black_box({
                    map.contains_key(*s);

                });
            }
        });
    }

        #[bench]
    fn u64_insert_btree(b: &mut Bencher) {
        let data = get_random_range(VEC_COUNT);

        let mut map = BTreeMap::new();

        b.iter(|| {
            map.clear();
            for s in data.iter() {
                test::black_box(map.insert(s, s));
            }
        });
    }

    #[bench]
    fn u64_get_btree(b: &mut Bencher) {
        let data = get_random_range(VEC_COUNT);

        let mut map = BTreeMap::new();
        for s in data.iter() {
            map.insert(s, s);
        }

        b.iter(|| {
            for s in data.iter() {
                test::black_box({
                    map.contains_key(s);

                });
            }
        });
    }

    fn get_random_range(count: usize) -> Vec<u64> {
        use rand::{Rng, SeedableRng, StdRng};

        let mut vec = Vec::new();

        let seed: &[_] = &[4, 2, 4, 2];
        let mut rng: StdRng = SeedableRng::from_seed(seed);

        for _ in 0..count {
            vec.push(rng.gen::<u64>());
        }

        vec.sort();
        vec.dedup();

        vec
    }
}
<commit_msg>Updated benchmarks<commit_after>#![feature(test)]

extern crate intmap;
extern crate rand;
extern crate test;

use intmap::IntMap;
use std::collections::HashMap;

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

    const VEC_COUNT: usize = 1000;

    #[bench]
    fn u64_insert_built_in(b: &mut Bencher) {
        let data = get_random_range(VEC_COUNT);
        let mut map = HashMap::with_capacity(data.len());

        b.iter(|| {
            map.clear();

            for s in data.iter() {
               test::black_box(map.insert(s, s));
            }

        });
    }

    #[bench]
    fn u64_get_built_in(b: &mut Bencher) {
        let data = get_random_range(VEC_COUNT);
        let mut map: HashMap<&u64, &u64>  = HashMap::with_capacity(data.len());

        for s in data.iter() {
            test::black_box(map.insert(s, s)
            );
        }

        b.iter(|| {
            for s in data.iter() {
                test::black_box({
                    map.contains_key(s);

                });
            }
        });
    }

    #[bench]
    fn u64_insert_intmap(b: &mut Bencher) {
        let data = get_random_range(VEC_COUNT);

        let mut map = IntMap::with_capacity(data.len());

        b.iter(|| {
            map.clear();
            for s in data.iter() {
                test::black_box({map.insert(*s, s)});
            }
        });
    }

    #[bench]
    fn u64_get_intmap(b: &mut Bencher) {
        let data = get_random_range(VEC_COUNT);

        let mut map = IntMap::with_capacity(data.len());
        for s in data.iter() {
            map.insert(*s, s);
        }

        b.iter(|| {
            for s in data.iter() {
                test::black_box(map.contains_key(*s));
            }
        });
    }

    fn get_random_range(count: usize) -> Vec<u64> {
        use rand::{Rng, SeedableRng, StdRng};

        let mut vec = Vec::new();

        let seed: &[_] = &[4, 2, 4, 2];
        let mut rng: StdRng = SeedableRng::from_seed(seed);

        for _ in 0..count {
            vec.push(rng.gen::<u64>());
        }

        vec.sort();
        vec.dedup();

        vec
    }
}
<|endoftext|>