<commit_before>// This is free and unencumbered software released into the public domain.
//
// Anyone is free to copy, modify, publish, use, compile, sell, or
// distribute this software, either in source code form or as a compiled
// binary, for any purpose, commercial or non-commercial, and by any
// means.
//
// In jurisdictions that recognize copyright laws, the author or authors
// of this software dedicate any and all copyright interest in the
// software to the public domain. We make this dedication for the benefit
// of the public at large and to the detriment of our heirs and
// successors. We intend this dedication to be an overt act of
// relinquishment in perpetuity of all present and future rights to this
// software under copyright law.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
// EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
// IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
// OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
// ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
// OTHER DEALINGS IN THE SOFTWARE.
//
// For more information, please refer to <http://unlicense.org/>

// This line shouldn't be necessary in a real project.
extern crate hyper_pkg as hyper;

use std::fs::File;
use std::path::PathBuf;

use futures::prelude::*;
use http::header::{HeaderValue, AUTHORIZATION, CONTENT_TYPE};
use serde::de;
use serde::Deserialize;
use twitter_stream::Token;

#[derive(Deserialize)]
#[serde(untagged)]
enum StreamMessage {
    Tweet(Tweet),
    Other(de::IgnoredAny),
}

struct Tweet {
    entities: Option<Entities>,
    id: u64,
    text: String,
    user: User,
    is_retweet: bool,
}

#[derive(Deserialize)]
struct Entities {
    user_mentions: Vec<UserMention>,
}

#[derive(Deserialize)]
struct UserMention {
    id: u64,
}

#[derive(Deserialize)]
struct User {
    id: u64,
    screen_name: String,
}

/// Represents a GET statuses/update request.
#[derive(oauth::Request)]
struct StatusUpdate<'a> {
    status: &'a str,
    in_reply_to_status_id: Option<u64>,
}

#[derive(Deserialize)]
#[serde(remote = "Token")]
struct TokenDef {
    // The `getter` attribute is required to make the `Deserialize` impl use the `From` conversion,
    // even if we are not deriving `Serialize` here.
    #[serde(getter = "__")]
    consumer_key: String,
    consumer_secret: String,
    access_key: String,
    access_secret: String,
}

impl From<TokenDef> for Token {
    fn from(def: TokenDef) -> Token {
        Token::from_parts(
            def.consumer_key,
            def.consumer_secret,
            def.access_key,
            def.access_secret,
        )
    }
}

type HttpsConnector = hyper_tls::HttpsConnector<hyper::client::HttpConnector>;

#[tokio::main]
async fn main() {
    let mut credential_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
    credential_path.push("credential.json");

    let conn = HttpsConnector::new();
    let mut client = hyper::Client::builder().build::<_, hyper::Body>(conn);

    let credential = File::open(credential_path).unwrap();
    let token =
        TokenDef::deserialize(&mut serde_json::Deserializer::from_reader(credential)).unwrap();

    // Information of the authenticated user:
    let user = verify_credentials(&token, &client).await;

    let mut stream = twitter_stream::Builder::new(token.as_ref())
        .track(format!("@{}", user.screen_name))
        .listen_with_client(&mut client)
        .try_flatten_stream();

    while let Some(json) = stream.next().await {
        if let Ok(StreamMessage::Tweet(tweet)) = serde_json::from_str(&json.unwrap()) {
            let mentions_me = |entities: Entities| {
                entities
                    .user_mentions
                    .iter()
                    .any(|mention| mention.id == user.id)
            };
            if !tweet.is_retweet
                && tweet.user.id != user.id
                && tweet.entities.map_or(false, mentions_me)
            {
                // Send a reply
                let tweeting = StatusUpdate {
                    status: &format!("@{} {}", tweet.user.screen_name, tweet.text),
                    in_reply_to_status_id: Some(tweet.id),
                }
                .send(&token, &client);
                tokio::spawn(tweeting);
            }
        }
    }
}

// The custom `Deserialize` impl is needed to handle Tweets with >140 characters.
// Once the Labs API is available, this impl should be unnecessary.
impl<'de> Deserialize<'de> for Tweet {
    fn deserialize<D: de::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
        #[derive(Deserialize)]
        struct Prototype {
            entities: Option<Entities>,
            id: u64,
            // Add the following attribute if you want to deserialize REST API responses too.
            // #[serde(alias = "full_text")]
            text: String,
            user: User,
            extended_tweet: Option<ExtendedTweet>,
            retweeted_status: Option<de::IgnoredAny>,
        }

        #[derive(Deserialize)]
        struct ExtendedTweet {
            full_text: String,
            entities: Option<Entities>,
        }

        Prototype::deserialize(d).map(|p| {
            let (text, entities) = p
                .extended_tweet
                .map_or((p.text, p.entities), |e| (e.full_text, e.entities));
            Tweet {
                entities,
                id: p.id,
                text,
                user: p.user,
                is_retweet: p.retweeted_status.is_some(),
            }
        })
    }
}

impl<'a> StatusUpdate<'a> {
    /// Performs the GET statuses/update request.
    fn send(
        &self,
        token: &Token,
        client: &hyper::Client<HttpsConnector>,
    ) -> impl Future<Output = ()> {
        const URI: &str = "https://api.twitter.com/1.1/statuses/update.json";

        let authorization = oauth::post(URI, self, token, oauth::HmacSha1);
        let form = oauth::to_form_urlencoded(self);
        let req = http::Request::post(http::Uri::from_static(URI))
            .header(
                CONTENT_TYPE,
                HeaderValue::from_static("application/x-www-form-urlencoded"),
            )
            .header(AUTHORIZATION, authorization)
            .body(form.into())
            .unwrap();

        parse_response::<de::IgnoredAny>(client.request(req)).map(|_| ())
    }
}

/// Performs a GET account/verify_credentials request.
async fn verify_credentials(token: &Token, client: &hyper::Client<HttpsConnector>) -> User {
    const URI: &str = "https://api.twitter.com/1.1/account/verify_credentials.json";

    let authorization = oauth::get(URI, &(), token, oauth::HmacSha1);
    let req = http::Request::get(http::Uri::from_static(URI))
        .header(AUTHORIZATION, authorization)
        .body(Default::default())
        .unwrap();

    parse_response(client.request(req)).await
}

fn parse_response<T: de::DeserializeOwned>(
    res: hyper::client::ResponseFuture,
) -> impl Future<Output = T> {
    res.then(|res| {
        let res = res.unwrap();
        if !res.status().is_success() {
            panic!("HTTP error: {}", res.status());
        }
        hyper::body::to_bytes(res).map(|body| serde_json::from_slice(&body.unwrap()).unwrap())
    })
}
<commit_msg>example: Rename `Prototype` to `Tweet`<commit_after>// This is free and unencumbered software released into the public domain.
//
// Anyone is free to copy, modify, publish, use, compile, sell, or
// distribute this software, either in source code form or as a compiled
// binary, for any purpose, commercial or non-commercial, and by any
// means.
//
// In jurisdictions that recognize copyright laws, the author or authors
// of this software dedicate any and all copyright interest in the
// software to the public domain. We make this dedication for the benefit
// of the public at large and to the detriment of our heirs and
// successors. We intend this dedication to be an overt act of
// relinquishment in perpetuity of all present and future rights to this
// software under copyright law.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
// EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
// IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
// OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
// ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
// OTHER DEALINGS IN THE SOFTWARE.
//
// For more information, please refer to <http://unlicense.org/>

// This line shouldn't be necessary in a real project.
extern crate hyper_pkg as hyper;

use std::fs::File;
use std::path::PathBuf;

use futures::prelude::*;
use http::header::{HeaderValue, AUTHORIZATION, CONTENT_TYPE};
use serde::de;
use serde::Deserialize;
use twitter_stream::Token;

#[derive(Deserialize)]
#[serde(untagged)]
enum StreamMessage {
    Tweet(Tweet),
    Other(de::IgnoredAny),
}

struct Tweet {
    entities: Option<Entities>,
    id: u64,
    text: String,
    user: User,
    is_retweet: bool,
}

#[derive(Deserialize)]
struct Entities {
    user_mentions: Vec<UserMention>,
}

#[derive(Deserialize)]
struct UserMention {
    id: u64,
}

#[derive(Deserialize)]
struct User {
    id: u64,
    screen_name: String,
}

/// Represents a GET statuses/update request.
#[derive(oauth::Request)]
struct StatusUpdate<'a> {
    status: &'a str,
    in_reply_to_status_id: Option<u64>,
}

#[derive(Deserialize)]
#[serde(remote = "Token")]
struct TokenDef {
    // The `getter` attribute is required to make the `Deserialize` impl use the `From` conversion,
    // even if we are not deriving `Serialize` here.
    #[serde(getter = "__")]
    consumer_key: String,
    consumer_secret: String,
    access_key: String,
    access_secret: String,
}

impl From<TokenDef> for Token {
    fn from(def: TokenDef) -> Token {
        Token::from_parts(
            def.consumer_key,
            def.consumer_secret,
            def.access_key,
            def.access_secret,
        )
    }
}

type HttpsConnector = hyper_tls::HttpsConnector<hyper::client::HttpConnector>;

#[tokio::main]
async fn main() {
    let mut credential_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
    credential_path.push("credential.json");

    let conn = HttpsConnector::new();
    let mut client = hyper::Client::builder().build::<_, hyper::Body>(conn);

    let credential = File::open(credential_path).unwrap();
    let token =
        TokenDef::deserialize(&mut serde_json::Deserializer::from_reader(credential)).unwrap();

    // Information of the authenticated user:
    let user = verify_credentials(&token, &client).await;

    let mut stream = twitter_stream::Builder::new(token.as_ref())
        .track(format!("@{}", user.screen_name))
        .listen_with_client(&mut client)
        .try_flatten_stream();

    while let Some(json) = stream.next().await {
        if let Ok(StreamMessage::Tweet(tweet)) = serde_json::from_str(&json.unwrap()) {
            let mentions_me = |entities: Entities| {
                entities
                    .user_mentions
                    .iter()
                    .any(|mention| mention.id == user.id)
            };
            if !tweet.is_retweet
                && tweet.user.id != user.id
                && tweet.entities.map_or(false, mentions_me)
            {
                // Send a reply
                let tweeting = StatusUpdate {
                    status: &format!("@{} {}", tweet.user.screen_name, tweet.text),
                    in_reply_to_status_id: Some(tweet.id),
                }
                .send(&token, &client);
                tokio::spawn(tweeting);
            }
        }
    }
}

// The custom `Deserialize` impl is needed to handle Tweets with >140 characters.
// Once the Labs API is available, this impl should be unnecessary.
impl<'de> Deserialize<'de> for Tweet {
    fn deserialize<D: de::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
        #[derive(Deserialize)]
        struct Tweet {
            entities: Option<Entities>,
            id: u64,
            // Add the following attribute if you want to deserialize REST API responses too.
            // #[serde(alias = "full_text")]
            text: String,
            user: User,
            extended_tweet: Option<ExtendedTweet>,
            retweeted_status: Option<de::IgnoredAny>,
        }

        #[derive(Deserialize)]
        struct ExtendedTweet {
            full_text: String,
            entities: Option<Entities>,
        }

        Tweet::deserialize(d).map(|t| {
            let (text, entities) = t
                .extended_tweet
                .map_or((t.text, t.entities), |e| (e.full_text, e.entities));
            Self {
                entities,
                id: t.id,
                text,
                user: t.user,
                is_retweet: t.retweeted_status.is_some(),
            }
        })
    }
}

impl<'a> StatusUpdate<'a> {
    /// Performs the GET statuses/update request.
    fn send(
        &self,
        token: &Token,
        client: &hyper::Client<HttpsConnector>,
    ) -> impl Future<Output = ()> {
        const URI: &str = "https://api.twitter.com/1.1/statuses/update.json";

        let authorization = oauth::post(URI, self, token, oauth::HmacSha1);
        let form = oauth::to_form_urlencoded(self);
        let req = http::Request::post(http::Uri::from_static(URI))
            .header(
                CONTENT_TYPE,
                HeaderValue::from_static("application/x-www-form-urlencoded"),
            )
            .header(AUTHORIZATION, authorization)
            .body(form.into())
            .unwrap();

        parse_response::<de::IgnoredAny>(client.request(req)).map(|_| ())
    }
}

/// Performs a GET account/verify_credentials request.
async fn verify_credentials(token: &Token, client: &hyper::Client<HttpsConnector>) -> User {
    const URI: &str = "https://api.twitter.com/1.1/account/verify_credentials.json";

    let authorization = oauth::get(URI, &(), token, oauth::HmacSha1);
    let req = http::Request::get(http::Uri::from_static(URI))
        .header(AUTHORIZATION, authorization)
        .body(Default::default())
        .unwrap();

    parse_response(client.request(req)).await
}

fn parse_response<T: de::DeserializeOwned>(
    res: hyper::client::ResponseFuture,
) -> impl Future<Output = T> {
    res.then(|res| {
        let res = res.unwrap();
        if !res.status().is_success() {
            panic!("HTTP error: {}", res.status());
        }
        hyper::body::to_bytes(res).map(|body| serde_json::from_slice(&body.unwrap()).unwrap())
    })
}
<|endoftext|>


<commit_before>/*
 * Cymbalum, Molecular Simulation in Rust
 * Copyright (C) 2015 Guillaume Fraux
 *
 * 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::f64::consts::PI;

use ::types::Vector3D;
use ::types::Matrix3;

/// The type of a cell determine how we will be able to compute the periodic
/// boundaries condition.
#[derive(Clone, Copy)]
pub enum CellType {
    /// Infinite unit cell, with no boundaries
    INFINITE,
    /// Orthorombic unit cell, with cuboide shape
    ORTHOROMBIC,
    /// Triclinic unit cell, with arbitrary parallelepipedic shape
    TRICLINIC,
}

/// The Universe type hold all the data about a system.
#[derive(Clone, Copy)]
pub struct UnitCell {
    data: Matrix3,
    celltype: CellType,
}

impl UnitCell {
    /// Create an infinite unit cell
    pub fn new() -> UnitCell {
        UnitCell{data: Matrix3::zero(), celltype: CellType::INFINITE}
    }
    /// Create an orthorombic unit cell
    pub fn ortho(a: f64, b: f64, c: f64) -> UnitCell {
        UnitCell{data: Matrix3::new(a, 0.0, 0.0,
                                    0.0, b, 0.0,
                                    0.0, 0.0, c),
                 celltype: CellType::ORTHOROMBIC}
    }
    /// Create a cubic unit cell
    pub fn cubic(L: f64) -> UnitCell {
        UnitCell{data: Matrix3::new(L, 0.0, 0.0,
                                    0.0, L, 0.0,
                                    0.0, 0.0, L),
                 celltype: CellType::ORTHOROMBIC}
    }
    /// Create a triclinic unit cell
    pub fn triclinic(a: f64, b: f64, c: f64, alpha: f64, beta: f64, gamma: f64) -> UnitCell {
        let cos_alpha = deg2rad(alpha).cos();
        let cos_beta = deg2rad(beta).cos();
        let (sin_gamma, cos_gamma) = deg2rad(gamma).sin_cos();

        let b_x = b * cos_gamma;
        let b_y = b * sin_gamma;

        let c_x = c * cos_beta;
        let c_y = c * (cos_alpha - cos_beta*cos_gamma)/sin_gamma;
        let c_z = f64::sqrt(c*c - c_y*c_y - c_x*c_x);

        UnitCell{data: Matrix3::new(a,   b_x, c_x,
                                    0.0, b_y, c_y,
                                    0.0, 0.0, c_z),
                 celltype: CellType::TRICLINIC}
    }

    /// Get the cell type
    pub fn celltype(&self) -> CellType {
        self.celltype
    }
    /// Set the cell type
    pub fn set_celltype(&mut self, ctype: CellType) {
        self.celltype = ctype;
    }

    /// Get the first vector of the cell
    pub fn vect_a(&self) -> Vector3D {
        let x = self.data[(0, 0)];
        let y = self.data[(1, 0)];
        let z = self.data[(2, 0)];
        Vector3D::new(x, y, z)
    }
    /// Get the first length of the cell
    pub fn a(&self) -> f64 {
        match self.celltype {
            CellType::TRICLINIC => self.vect_a().norm(),
            _ => self.data[(0, 0)],
        }
    }

    /// Get the second length of the cell
    pub fn vect_b(&self) -> Vector3D {
        let x = self.data[(0, 1)];
        let y = self.data[(1, 1)];
        let z = self.data[(2, 1)];
        Vector3D::new(x, y, z)
    }
    /// Get the second length of the cell
    pub fn b(&self) -> f64 {
        match self.celltype {
            CellType::TRICLINIC => self.vect_b().norm(),
            _ => self.data[(1, 1)],
        }
    }

    /// Get the third length of the cell
    pub fn vect_c(&self) -> Vector3D {
        let x = self.data[(0, 2)];
        let y = self.data[(1, 2)];
        let z = self.data[(2, 2)];
        Vector3D::new(x, y, z)
    }
    /// Get the second length of the cell
    pub fn c(&self) -> f64 {
        match self.celltype {
            CellType::TRICLINIC => self.vect_c().norm(),
            _ => self.data[(2, 2)],
        }
    }

    /// Get the first angle of the cell
    pub fn alpha(&self) -> f64 {
        match self.celltype {
            CellType::TRICLINIC => {
                let b = self.vect_b();
                let c = self.vect_c();
                rad2deg(angle(b, c))
            },
            _ => 90.0,
        }
    }

    /// Get the second angle of the cell
    pub fn beta(&self) -> f64 {
        match self.celltype {
            CellType::TRICLINIC => {
                let a = self.vect_a();
                let c = self.vect_c();
                rad2deg(angle(a, c))
            },
            _ => 90.0,
        }
    }

    /// Get the third angle of the cell
    pub fn gamma(&self) -> f64 {
        match self.celltype {
            CellType::TRICLINIC => {
                let a = self.vect_a();
                let b = self.vect_b();
                rad2deg(angle(a, b))
            },
            _ => 90.0,
        }
    }

    /// Get the volume angle of the cell
    pub fn volume(&self) -> f64 {
        match self.celltype {
            CellType::INFINITE => 0.0,
            CellType::ORTHOROMBIC => self.a()*self.b()*self.c(),
            CellType::TRICLINIC => {
                // The volume is the mixed product of the three cell vectors
                let a = self.vect_a();
                let b = self.vect_b();
                let c = self.vect_c();
                a * (b ^ c)
            },
        }
    }

    /// Scale this unit cell in-place by multiplying the H matrix by `s`.
    pub fn scale_mut(&mut self, s: Matrix3) {
        self.data = s * self.data;
    }

    /// Scale this unit cell by multiplying the H matrix by `s`, and return a new unit cell
    pub fn scale(&self, s: Matrix3) -> UnitCell {
        UnitCell{data: s * self.data, celltype: self.celltype}
    }

    /// Periodic boundary conditions distance between the point `u` and the point `v`
    pub fn distance(&self, u: &Vector3D, v: &Vector3D) -> f64 {
        let d = *v - *u;
        match self.celltype {
            CellType::INFINITE => d.norm(),
            CellType::ORTHOROMBIC => {
                let x = d.x - f64::round(d.x/self.a())*self.a();
                let y = d.y - f64::round(d.y/self.b())*self.b();
                let z = d.z - f64::round(d.z/self.c())*self.c();
                return f64::sqrt(x*x + y*y + z*z);
            },
            CellType::TRICLINIC => {
                let inv = self.data.inverse();
                let s = inv * d;
                let x = s.x - f64::round(s.x);
                let y = s.y - f64::round(s.y);
                let z = s.z - f64::round(s.z);
                let res = Vector3D::new(x, y, z);
                return (self.data*res).norm();
            },
        }
    }
}

/// Convert `x` from degrees to radians
fn deg2rad(x: f64) -> f64 {
    x * PI / 180.0
}
/// Convert `x` from radians to degrees
fn rad2deg(x: f64) -> f64 {
    x * 180.0 / PI
}

/// Get the angles between the vectors `u` and `v`.
fn angle(u: Vector3D, v: Vector3D) -> f64 {
    let un = u.normalized();
    let vn = v.normalized();
    f64::acos(un*vn)
}

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

    #[test]
    fn infinite() {
        let cell = UnitCell::new();
        assert_eq!(cell.a(), 0.0);
        assert_eq!(cell.b(), 0.0);
        assert_eq!(cell.c(), 0.0);

        assert_eq!(cell.alpha(), 90.0);
        assert_eq!(cell.beta(), 90.0);
        assert_eq!(cell.gamma(), 90.0);

        assert_eq!(cell.volume(), 0.0);
    }

    #[test]
    fn cubic() {
        let cell = UnitCell::cubic(3.0);
        assert_eq!(cell.a(), 3.0);
        assert_eq!(cell.b(), 3.0);
        assert_eq!(cell.c(), 3.0);

        assert_eq!(cell.alpha(), 90.0);
        assert_eq!(cell.beta(), 90.0);
        assert_eq!(cell.gamma(), 90.0);

        assert_eq!(cell.volume(), 3.0*3.0*3.0);
    }

    #[test]
    fn orthorombic() {
        let cell = UnitCell::ortho(3.0, 4.0, 5.0);
        assert_eq!(cell.a(), 3.0);
        assert_eq!(cell.b(), 4.0);
        assert_eq!(cell.c(), 5.0);

        assert_eq!(cell.alpha(), 90.0);
        assert_eq!(cell.beta(), 90.0);
        assert_eq!(cell.gamma(), 90.0);

        assert_eq!(cell.volume(), 3.0*4.0*5.0);
    }

    #[test]
    fn triclinic() {
        let cell = UnitCell::triclinic(3.0, 4.0, 5.0, 80.0, 90.0, 110.0);
        assert_eq!(cell.a(), 3.0);
        assert_eq!(cell.b(), 4.0);
        assert_eq!(cell.c(), 5.0);

        assert_eq!(cell.alpha(), 80.0);
        assert_eq!(cell.beta(), 90.0);
        assert_eq!(cell.gamma(), 110.0);

        assert_approx_eq!(cell.volume(), 55.410529, 1e-6);
    }

    #[test]
    fn scale() {
        let cell = UnitCell::ortho(3.0, 4.0, 5.0);
        let A = 2.0 * Matrix3::one();
        let cell = cell.scale(A);

        assert_eq!(cell.a(), 6.0);
        assert_eq!(cell.b(), 8.0);
        assert_eq!(cell.c(), 10.0);
    }

    #[test]
    fn scale_mut() {
        let mut cell = UnitCell::ortho(3.0, 4.0, 5.0);
        let A = 2.0 * Matrix3::one();
        cell.scale_mut(A);

        assert_eq!(cell.a(), 6.0);
        assert_eq!(cell.b(), 8.0);
        assert_eq!(cell.c(), 10.0);
    }

    #[test]
    fn distances() {
        // Orthorombic unit cell
        let cell = UnitCell::ortho(3.0, 4.0, 5.0);
        let u = &Vector3D::new(0.0, 0.0, 0.0);
        let v = &Vector3D::new(1.0, 2.0, 6.0);
        assert_eq!(cell.distance(u, v), f64::sqrt(6.0));

        // Infinite unit cell
        let cell = UnitCell::new();
        assert_eq!(cell.distance(u, v), v.norm());

        // Triclinic unit cell
        let cell = UnitCell::triclinic(3.0, 4.0, 5.0, 90.0, 90.0, 90.0);
        assert_eq!(cell.distance(u, v), f64::sqrt(6.0));
    }
}
<commit_msg>Add a wrap_vector function to UnitCell<commit_after>/*
 * Cymbalum, Molecular Simulation in Rust
 * Copyright (C) 2015 Guillaume Fraux
 *
 * 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::f64::consts::PI;

use ::types::Vector3D;
use ::types::Matrix3;

/// The type of a cell determine how we will be able to compute the periodic
/// boundaries condition.
#[derive(Clone, Copy)]
pub enum CellType {
    /// Infinite unit cell, with no boundaries
    INFINITE,
    /// Orthorombic unit cell, with cuboide shape
    ORTHOROMBIC,
    /// Triclinic unit cell, with arbitrary parallelepipedic shape
    TRICLINIC,
}

/// The Universe type hold all the data about a system.
#[derive(Clone, Copy)]
pub struct UnitCell {
    data: Matrix3,
    celltype: CellType,
}

impl UnitCell {
    /// Create an infinite unit cell
    pub fn new() -> UnitCell {
        UnitCell{data: Matrix3::zero(), celltype: CellType::INFINITE}
    }
    /// Create an orthorombic unit cell
    pub fn ortho(a: f64, b: f64, c: f64) -> UnitCell {
        UnitCell{data: Matrix3::new(a, 0.0, 0.0,
                                    0.0, b, 0.0,
                                    0.0, 0.0, c),
                 celltype: CellType::ORTHOROMBIC}
    }
    /// Create a cubic unit cell
    pub fn cubic(L: f64) -> UnitCell {
        UnitCell{data: Matrix3::new(L, 0.0, 0.0,
                                    0.0, L, 0.0,
                                    0.0, 0.0, L),
                 celltype: CellType::ORTHOROMBIC}
    }
    /// Create a triclinic unit cell
    pub fn triclinic(a: f64, b: f64, c: f64, alpha: f64, beta: f64, gamma: f64) -> UnitCell {
        let cos_alpha = deg2rad(alpha).cos();
        let cos_beta = deg2rad(beta).cos();
        let (sin_gamma, cos_gamma) = deg2rad(gamma).sin_cos();

        let b_x = b * cos_gamma;
        let b_y = b * sin_gamma;

        let c_x = c * cos_beta;
        let c_y = c * (cos_alpha - cos_beta*cos_gamma)/sin_gamma;
        let c_z = f64::sqrt(c*c - c_y*c_y - c_x*c_x);

        UnitCell{data: Matrix3::new(a,   b_x, c_x,
                                    0.0, b_y, c_y,
                                    0.0, 0.0, c_z),
                 celltype: CellType::TRICLINIC}
    }

    /// Get the cell type
    pub fn celltype(&self) -> CellType {
        self.celltype
    }
    /// Set the cell type
    pub fn set_celltype(&mut self, ctype: CellType) {
        self.celltype = ctype;
    }

    /// Get the first vector of the cell
    pub fn vect_a(&self) -> Vector3D {
        let x = self.data[(0, 0)];
        let y = self.data[(1, 0)];
        let z = self.data[(2, 0)];
        Vector3D::new(x, y, z)
    }
    /// Get the first length of the cell
    pub fn a(&self) -> f64 {
        match self.celltype {
            CellType::TRICLINIC => self.vect_a().norm(),
            _ => self.data[(0, 0)],
        }
    }

    /// Get the second length of the cell
    pub fn vect_b(&self) -> Vector3D {
        let x = self.data[(0, 1)];
        let y = self.data[(1, 1)];
        let z = self.data[(2, 1)];
        Vector3D::new(x, y, z)
    }
    /// Get the second length of the cell
    pub fn b(&self) -> f64 {
        match self.celltype {
            CellType::TRICLINIC => self.vect_b().norm(),
            _ => self.data[(1, 1)],
        }
    }

    /// Get the third length of the cell
    pub fn vect_c(&self) -> Vector3D {
        let x = self.data[(0, 2)];
        let y = self.data[(1, 2)];
        let z = self.data[(2, 2)];
        Vector3D::new(x, y, z)
    }
    /// Get the second length of the cell
    pub fn c(&self) -> f64 {
        match self.celltype {
            CellType::TRICLINIC => self.vect_c().norm(),
            _ => self.data[(2, 2)],
        }
    }

    /// Get the first angle of the cell
    pub fn alpha(&self) -> f64 {
        match self.celltype {
            CellType::TRICLINIC => {
                let b = self.vect_b();
                let c = self.vect_c();
                rad2deg(angle(b, c))
            },
            _ => 90.0,
        }
    }

    /// Get the second angle of the cell
    pub fn beta(&self) -> f64 {
        match self.celltype {
            CellType::TRICLINIC => {
                let a = self.vect_a();
                let c = self.vect_c();
                rad2deg(angle(a, c))
            },
            _ => 90.0,
        }
    }

    /// Get the third angle of the cell
    pub fn gamma(&self) -> f64 {
        match self.celltype {
            CellType::TRICLINIC => {
                let a = self.vect_a();
                let b = self.vect_b();
                rad2deg(angle(a, b))
            },
            _ => 90.0,
        }
    }

    /// Get the volume angle of the cell
    pub fn volume(&self) -> f64 {
        match self.celltype {
            CellType::INFINITE => 0.0,
            CellType::ORTHOROMBIC => self.a()*self.b()*self.c(),
            CellType::TRICLINIC => {
                // The volume is the mixed product of the three cell vectors
                let a = self.vect_a();
                let b = self.vect_b();
                let c = self.vect_c();
                a * (b ^ c)
            },
        }
    }

    /// Scale this unit cell in-place by multiplying the H matrix by `s`.
    pub fn scale_mut(&mut self, s: Matrix3) {
        self.data = s * self.data;
    }

    /// Scale this unit cell by multiplying the H matrix by `s`, and return a new unit cell
    pub fn scale(&self, s: Matrix3) -> UnitCell {
        UnitCell{data: s * self.data, celltype: self.celltype}
    }

    /// Wrap a vector in the unit cell, obeying the periodic boundary conditions
    pub fn wrap_vector(&self, vect: &mut Vector3D) {
        match self.celltype {
            CellType::INFINITE => (),
            CellType::ORTHOROMBIC => {
                vect.x = vect.x - f64::round(vect.x/self.a())*self.a();
                vect.y = vect.y - f64::round(vect.y/self.b())*self.b();
                vect.z = vect.z - f64::round(vect.z/self.c())*self.c();
            },
            CellType::TRICLINIC => {
                let inv = self.data.inverse();
                let mut fractional = inv * (*vect);
                fractional.x = fractional.x - f64::round(fractional.x);
                fractional.y = fractional.y - f64::round(fractional.y);
                fractional.z = fractional.z - f64::round(fractional.z);
                let res = self.data * fractional;
                vect.x = res.x;
                vect.y = res.y;
                vect.z = res.z;
            },
        }
    }

    /// Periodic boundary conditions distance between the point `u` and the point `v`
    pub fn distance(&self, u: &Vector3D, v: &Vector3D) -> f64 {
        let mut d = *v - *u;
        self.wrap_vector(&mut d);
        return d.norm();
    }
}

/// Convert `x` from degrees to radians
fn deg2rad(x: f64) -> f64 {
    x * PI / 180.0
}
/// Convert `x` from radians to degrees
fn rad2deg(x: f64) -> f64 {
    x * 180.0 / PI
}

/// Get the angles between the vectors `u` and `v`.
fn angle(u: Vector3D, v: Vector3D) -> f64 {
    let un = u.normalized();
    let vn = v.normalized();
    f64::acos(un*vn)
}

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

    #[test]
    fn infinite() {
        let cell = UnitCell::new();
        assert_eq!(cell.a(), 0.0);
        assert_eq!(cell.b(), 0.0);
        assert_eq!(cell.c(), 0.0);

        assert_eq!(cell.alpha(), 90.0);
        assert_eq!(cell.beta(), 90.0);
        assert_eq!(cell.gamma(), 90.0);

        assert_eq!(cell.volume(), 0.0);
    }

    #[test]
    fn cubic() {
        let cell = UnitCell::cubic(3.0);
        assert_eq!(cell.a(), 3.0);
        assert_eq!(cell.b(), 3.0);
        assert_eq!(cell.c(), 3.0);

        assert_eq!(cell.alpha(), 90.0);
        assert_eq!(cell.beta(), 90.0);
        assert_eq!(cell.gamma(), 90.0);

        assert_eq!(cell.volume(), 3.0*3.0*3.0);
    }

    #[test]
    fn orthorombic() {
        let cell = UnitCell::ortho(3.0, 4.0, 5.0);
        assert_eq!(cell.a(), 3.0);
        assert_eq!(cell.b(), 4.0);
        assert_eq!(cell.c(), 5.0);

        assert_eq!(cell.alpha(), 90.0);
        assert_eq!(cell.beta(), 90.0);
        assert_eq!(cell.gamma(), 90.0);

        assert_eq!(cell.volume(), 3.0*4.0*5.0);
    }

    #[test]
    fn triclinic() {
        let cell = UnitCell::triclinic(3.0, 4.0, 5.0, 80.0, 90.0, 110.0);
        assert_eq!(cell.a(), 3.0);
        assert_eq!(cell.b(), 4.0);
        assert_eq!(cell.c(), 5.0);

        assert_eq!(cell.alpha(), 80.0);
        assert_eq!(cell.beta(), 90.0);
        assert_eq!(cell.gamma(), 110.0);

        assert_approx_eq!(cell.volume(), 55.410529, 1e-6);
    }

    #[test]
    fn scale() {
        let cell = UnitCell::ortho(3.0, 4.0, 5.0);
        let A = 2.0 * Matrix3::one();
        let cell = cell.scale(A);

        assert_eq!(cell.a(), 6.0);
        assert_eq!(cell.b(), 8.0);
        assert_eq!(cell.c(), 10.0);
    }

    #[test]
    fn scale_mut() {
        let mut cell = UnitCell::ortho(3.0, 4.0, 5.0);
        let A = 2.0 * Matrix3::one();
        cell.scale_mut(A);

        assert_eq!(cell.a(), 6.0);
        assert_eq!(cell.b(), 8.0);
        assert_eq!(cell.c(), 10.0);
    }

    #[test]
    fn distances() {
        // Orthorombic unit cell
        let cell = UnitCell::ortho(3.0, 4.0, 5.0);
        let u = &Vector3D::new(0.0, 0.0, 0.0);
        let v = &Vector3D::new(1.0, 2.0, 6.0);
        assert_eq!(cell.distance(u, v), f64::sqrt(6.0));

        // Infinite unit cell
        let cell = UnitCell::new();
        assert_eq!(cell.distance(u, v), v.norm());

        // Triclinic unit cell
        let cell = UnitCell::triclinic(3.0, 4.0, 5.0, 90.0, 90.0, 90.0);
        assert_eq!(cell.distance(u, v), f64::sqrt(6.0));
    }
}
<|endoftext|>


<commit_before>//! An example that uses CPUID to determine the system topology.
//!
//! Intel Topology is a pretty complicated subject (unfortunately):
//! https://software.intel.com/en-us/articles/intel-64-architecture-processor-topology-enumeration/
extern crate core_affinity;
extern crate raw_cpuid;

use raw_cpuid::{CpuId, ExtendedTopologyLevel, TopologyType};
use std::thread;
use std::convert::TryInto;

/// Runs CPU ID on every core in the system (to gather all APIC IDs).
///
/// # Note
/// This won't work on macOS, apparently there is no guarantee after setting the affinity
/// that a thread will really execute on the pinned core.
fn gather_all_xapic_ids() -> Vec<u8> {
    let core_ids = core_affinity::get_core_ids().unwrap();

    // Create a thread for each active CPU core.
    core_ids
        .into_iter()
        .map(|id| {
            thread::spawn(move || {
                // Pin this thread to a single CPU core.
                core_affinity::set_for_current(id);
                // Do more work after this.
                let cpuid = CpuId::new();
                cpuid
                    .get_feature_info()
                    .map_or_else(|| 0, |finfo| finfo.initial_local_apic_id())
            })
            .join()
            .unwrap_or(0)
        })
        .collect::<Vec<_>>()
}

/// Runs CPU ID on every core in the system (to gather all x2APIC IDs).
///
/// # Note
/// This won't work on macOS, apparently there is no guarantee after setting the affinity
/// that a thread will really execute on the pinned core.
fn gather_all_x2apic_ids() -> Vec<u32> {
    let core_ids = core_affinity::get_core_ids().unwrap();

    // Create a thread for each active CPU core.
    core_ids
        .into_iter()
        .map(|id| {
            thread::spawn(move || {
                // Pin this thread to a single CPU core.
                core_affinity::set_for_current(id);
                // Do more work after this.
                let cpuid = CpuId::new();
                cpuid.get_extended_topology_info().map_or_else(
                    || 0,
                    |mut topiter| topiter.next().as_ref().unwrap().x2apic_id(),
                )
            })
            .join()
            .unwrap_or(0)
        })
        .collect::<Vec<_>>()
}

fn enumerate_with_x2apic_ids() {
    let cpuid = CpuId::new();
    let mut smt_x2apic_shift: u32 = 0;
    let mut core_x2apic_shift: u32 = 0;

    cpuid.get_extended_topology_info().map_or_else(
        || println!("No topology information available."),
        |topoiter| {
            let topology: Vec<ExtendedTopologyLevel> = topoiter.collect();
            for topolevel in topology.iter() {
                match topolevel.level_type() {
                    TopologyType::SMT => {
                        smt_x2apic_shift = topolevel.shift_right_for_next_apic_id();
                    }
                    TopologyType::Core => {
                        core_x2apic_shift = topolevel.shift_right_for_next_apic_id();
                    }
                    _ => panic!("Topology category not supported."),
                };
            }
        },
    );

    println!("Enumeration of all cores in the system (with x2APIC IDs):");
    let mut all_x2apic_ids: Vec<u32> = gather_all_x2apic_ids();
    all_x2apic_ids.sort_unstable();
    for x2apic_id in all_x2apic_ids {
        let smt_select_mask = !(u32::max_value() << smt_x2apic_shift);
        let core_select_mask = (!((u32::max_value()) << core_x2apic_shift)) ^ smt_select_mask;
        let pkg_select_mask = u32::max_value() << core_x2apic_shift;

        let smt_id = x2apic_id & smt_select_mask;
        let core_id = (x2apic_id & core_select_mask) >> smt_x2apic_shift;
        let pkg_id = (x2apic_id & pkg_select_mask) >> core_x2apic_shift;

        println!(
            "x2APIC#{} (pkg: {}, core: {}, smt: {})",
            x2apic_id, pkg_id, core_id, smt_id
        );
    }
}

fn cpuid_bits_needed(count: u8) -> u8 {
    let mut mask: u8 = 0x80;
    let mut cnt: u8 = 8;

    while (cnt > 0) && ((mask & count) != mask) {
        mask >>= 1;
        cnt -= 1;
    }

    cnt
}

fn get_processor_limits() -> (u8, u8) {
    let cpuid = CpuId::new();

    // This is for AMD processors:
    if let Some(info) = cpuid.get_processor_capacity_feature_info() {
        let max_logical_processor_ids = info.num_phys_threads();
        let smt_max_cores_for_package = info.apic_id_size();

        return (max_logical_processor_ids.try_into().unwrap(), smt_max_cores_for_package.try_into().unwrap());
    }
    // This is for Intel processors:
    else if let Some(cparams) = cpuid.get_cache_parameters() {
        let max_logical_processor_ids = cpuid
            .get_feature_info()
            .map_or_else(|| 1, |finfo| finfo.max_logical_processor_ids());

        let mut smt_max_cores_for_package: u8 = 1;
        for (ecx, cache) in cparams.enumerate() {
            if ecx == 0 {
                smt_max_cores_for_package = cache.max_cores_for_package() as u8;
            }
        }

        return (max_logical_processor_ids as u8, smt_max_cores_for_package as u8);
    }

    unreachable!("Example doesn't support this CPU")
}

fn enumerate_with_xapic_id() {
    let (max_logical_processor_ids, smt_max_cores_for_package) = get_processor_limits();

    let smt_mask_width: u8 = cpuid_bits_needed(
        (max_logical_processor_ids.next_power_of_two() / smt_max_cores_for_package) - 1,
    );
    let smt_select_mask: u8 = !(u8::max_value() << smt_mask_width);
    let core_mask_width: u8 = cpuid_bits_needed(smt_max_cores_for_package - 1);
    let core_only_select_mask =
        (!(u8::max_value() << (core_mask_width + smt_mask_width))) ^ smt_select_mask;
    let pkg_select_mask = u8::max_value() << (core_mask_width + smt_mask_width);

    println!("Enumeration of all cores in the system (with APIC IDs):");
    let mut all_xapic_ids: Vec<u8> = gather_all_xapic_ids();
    all_xapic_ids.sort_unstable();

    for xapic_id in all_xapic_ids {
        let smt_id = xapic_id & smt_select_mask;
        let core_id = (xapic_id & core_only_select_mask) >> smt_mask_width;
        let pkg_id = (xapic_id & pkg_select_mask) >> (core_mask_width + smt_mask_width);

        println!(
            "APIC#{} (pkg: {}, core: {}, smt: {})",
            xapic_id, pkg_id, core_id, smt_id
        );
    }
}

fn main() {
    let cpuid = CpuId::new();

    cpuid.get_processor_brand_string().map_or_else(
        || println!("CPU model identifier not available."),
        |pbs| println!("CPU Model is: {}", pbs.as_str()),
    );
    cpuid.get_extended_topology_info().map_or_else(
        || println!("No topology information available."),
        |topoiter| {
            let mut topology: Vec<ExtendedTopologyLevel> = topoiter.collect();
            topology.reverse();

            for topolevel in topology.iter() {
                let typ = match topolevel.level_type() {
                    TopologyType::SMT => "SMT-threads",
                    TopologyType::Core => "cores",
                    _ => panic!("Topology category not supported."),
                };

                println!(
                    "At level {} the CPU has: {} {}",
                    topolevel.level_number(),
                    topolevel.processors(),
                    typ,
                );
            }
        },
    );

    println!();
    enumerate_with_xapic_id();

    println!();
    enumerate_with_x2apic_ids();
}
<commit_msg>A word.<commit_after>//! An example that uses CPUID to determine the system topology.
//!
//! Intel Topology is a pretty complicated subject (unfortunately):
//! https://software.intel.com/en-us/articles/intel-64-architecture-processor-topology-enumeration/
extern crate core_affinity;
extern crate raw_cpuid;

use raw_cpuid::{CpuId, ExtendedTopologyLevel, TopologyType};
use std::thread;
use std::convert::TryInto;

/// Runs CPU ID on every core in the system (to gather all APIC IDs).
///
/// # Note
/// This won't work on macOS, apparently there is no guarantee after setting the affinity
/// that a thread will really execute on the pinned core.
fn gather_all_xapic_ids() -> Vec<u8> {
    let core_ids = core_affinity::get_core_ids().unwrap();

    // Create a thread for each active CPU core.
    core_ids
        .into_iter()
        .map(|id| {
            thread::spawn(move || {
                // Pin this thread to a single CPU core.
                core_affinity::set_for_current(id);
                // Do more work after this.
                let cpuid = CpuId::new();
                cpuid
                    .get_feature_info()
                    .map_or_else(|| 0, |finfo| finfo.initial_local_apic_id())
            })
            .join()
            .unwrap_or(0)
        })
        .collect::<Vec<_>>()
}

/// Runs CPU ID on every core in the system (to gather all x2APIC IDs).
///
/// # Note
/// This won't work on macOS, apparently there is no guarantee after setting the affinity
/// that a thread will really execute on the pinned core.
fn gather_all_x2apic_ids() -> Vec<u32> {
    let core_ids = core_affinity::get_core_ids().unwrap();

    // Create a thread for each active CPU core.
    core_ids
        .into_iter()
        .map(|id| {
            thread::spawn(move || {
                // Pin this thread to a single CPU core.
                core_affinity::set_for_current(id);
                // Do more work after this.
                let cpuid = CpuId::new();
                cpuid.get_extended_topology_info().map_or_else(
                    || 0,
                    |mut topiter| topiter.next().as_ref().unwrap().x2apic_id(),
                )
            })
            .join()
            .unwrap_or(0)
        })
        .collect::<Vec<_>>()
}

fn enumerate_with_x2apic_ids() {
    let cpuid = CpuId::new();
    let mut smt_x2apic_shift: u32 = 0;
    let mut core_x2apic_shift: u32 = 0;

    cpuid.get_extended_topology_info().map_or_else(
        || println!("No topology information available."),
        |topoiter| {
            let topology: Vec<ExtendedTopologyLevel> = topoiter.collect();
            for topolevel in topology.iter() {
                match topolevel.level_type() {
                    TopologyType::SMT => {
                        smt_x2apic_shift = topolevel.shift_right_for_next_apic_id();
                    }
                    TopologyType::Core => {
                        core_x2apic_shift = topolevel.shift_right_for_next_apic_id();
                    }
                    _ => panic!("Topology category not supported."),
                };
            }
        },
    );

    println!("Enumeration of all cores in the system (with x2APIC IDs):");
    let mut all_x2apic_ids: Vec<u32> = gather_all_x2apic_ids();
    all_x2apic_ids.sort_unstable();
    for x2apic_id in all_x2apic_ids {
        let smt_select_mask = !(u32::max_value() << smt_x2apic_shift);
        let core_select_mask = (!((u32::max_value()) << core_x2apic_shift)) ^ smt_select_mask;
        let pkg_select_mask = u32::max_value() << core_x2apic_shift;

        let smt_id = x2apic_id & smt_select_mask;
        let core_id = (x2apic_id & core_select_mask) >> smt_x2apic_shift;
        let pkg_id = (x2apic_id & pkg_select_mask) >> core_x2apic_shift;

        println!(
            "x2APIC#{} (pkg: {}, core: {}, smt: {})",
            x2apic_id, pkg_id, core_id, smt_id
        );
    }
}

fn cpuid_bits_needed(count: u8) -> u8 {
    let mut mask: u8 = 0x80;
    let mut cnt: u8 = 8;

    while (cnt > 0) && ((mask & count) != mask) {
        mask >>= 1;
        cnt -= 1;
    }

    cnt
}

fn get_processor_limits() -> (u8, u8) {
    let cpuid = CpuId::new();

    // This is for AMD processors:
    if let Some(info) = cpuid.get_processor_capacity_feature_info() {
        let max_logical_processor_ids = info.num_phys_threads();
        let smt_max_cores_for_package = info.apic_id_size();

        return (max_logical_processor_ids.try_into().unwrap(), smt_max_cores_for_package.try_into().unwrap());
    }
    // This is for Intel processors:
    else if let Some(cparams) = cpuid.get_cache_parameters() {
        let max_logical_processor_ids = cpuid
            .get_feature_info()
            .map_or_else(|| 1, |finfo| finfo.max_logical_processor_ids());

        let mut smt_max_cores_for_package: u8 = 1;
        for (ecx, cache) in cparams.enumerate() {
            if ecx == 0 {
                smt_max_cores_for_package = cache.max_cores_for_package() as u8;
            }
        }

        return (max_logical_processor_ids as u8, smt_max_cores_for_package as u8);
    }

    unreachable!("Example doesn't support this CPU")
}

fn enumerate_with_xapic_ids() {
    let (max_logical_processor_ids, smt_max_cores_for_package) = get_processor_limits();

    let smt_mask_width: u8 = cpuid_bits_needed(
        (max_logical_processor_ids.next_power_of_two() / smt_max_cores_for_package) - 1,
    );
    let smt_select_mask: u8 = !(u8::max_value() << smt_mask_width);
    let core_mask_width: u8 = cpuid_bits_needed(smt_max_cores_for_package - 1);
    let core_only_select_mask =
        (!(u8::max_value() << (core_mask_width + smt_mask_width))) ^ smt_select_mask;
    let pkg_select_mask = u8::max_value() << (core_mask_width + smt_mask_width);

    println!("Enumeration of all cores in the system (with APIC IDs):");
    let mut all_xapic_ids: Vec<u8> = gather_all_xapic_ids();
    all_xapic_ids.sort_unstable();

    for xapic_id in all_xapic_ids {
        let smt_id = xapic_id & smt_select_mask;
        let core_id = (xapic_id & core_only_select_mask) >> smt_mask_width;
        let pkg_id = (xapic_id & pkg_select_mask) >> (core_mask_width + smt_mask_width);

        println!(
            "APIC#{} (pkg: {}, core: {}, smt: {})",
            xapic_id, pkg_id, core_id, smt_id
        );
    }
}

fn main() {
    let cpuid = CpuId::new();

    cpuid.get_processor_brand_string().map_or_else(
        || println!("CPU model identifier not available."),
        |pbs| println!("CPU Model is: {}", pbs.as_str()),
    );
    cpuid.get_extended_topology_info().map_or_else(
        || println!("No topology information available."),
        |topoiter| {
            let mut topology: Vec<ExtendedTopologyLevel> = topoiter.collect();
            topology.reverse();

            for topolevel in topology.iter() {
                let typ = match topolevel.level_type() {
                    TopologyType::SMT => "SMT-threads",
                    TopologyType::Core => "cores",
                    _ => panic!("Topology category not supported."),
                };

                println!(
                    "At level {} the CPU has: {} {}",
                    topolevel.level_number(),
                    topolevel.processors(),
                    typ,
                );
            }
        },
    );

    println!();
    enumerate_with_xapic_ids();

    println!();
    enumerate_with_x2apic_ids();
}
<|endoftext|>


<commit_before>

use visit::IntoNeighbors;
use visit::{VisitMap, Visitable};

/// Strictly monotonically increasing event time for a depth first search.
#[derive(Copy, Clone, Debug, PartialEq, PartialOrd, Eq, Ord, Default)]
pub struct Time(pub usize);

/// A depth first search (DFS) visitor event.
#[derive(Copy, Clone, Debug)]
pub enum DfsEvent<N> {
    Discover(N, Time),
    /// An edge of the tree formed by the traversal.
    TreeEdge(N, N),
    /// An edge to an already visited node.
    BackEdge(N, N),
    /// A cross or forward edge.
    ///
    /// For an edge *(u, v)*, if the discover time of *v* is greater than *u*,
    /// then it is a forward edge, else a cross edge.
    CrossForwardEdge(N, N),
    Finish(N, Time),
}

/// Return if the expression is a break value.
macro_rules! try_control {
    ($e:expr) => {
        match $e {
            x => if x.should_break() {
                return x;
            }
        }
    }
}

/// Control flow for callbacks.
///
/// `Break` can carry a value.
#[derive(Copy, Clone, Debug)]
pub enum Control<B> {
    Continue,
    Break(B),
}

impl<B> Control<B> {
    pub fn breaking() -> Control<()> { Control::Break(()) }
    /// Get the value in `Control::Break(_)`, if present.
    pub fn break_value(self) -> Option<B> {
        match self {
            Control::Continue => None,
            Control::Break(b) => Some(b),
        }
    }
}

/// Control flow for callbacks.
///
/// The empty return value `()` is equivalent to continue.
pub trait ControlFlow {
    fn continuing() -> Self;
    #[inline]
    fn should_break(&self) -> bool { false }
}

impl ControlFlow for () {
    fn continuing() { }
}

impl<B> ControlFlow for Control<B> {
    fn continuing() -> Self { Control::Continue }
    fn should_break(&self) -> bool {
        if let Control::Break(_) = *self { true } else { false }
    }
}

impl<E> ControlFlow for Result<(), E> {
    fn continuing() -> Self { Ok(()) }
    fn should_break(&self) -> bool {
        if let Err(_) = *self { true } else { false }
    }
}

/// The default is `Continue`.
impl<B> Default for Control<B> {
    fn default() -> Self { Control::Continue }
}

/// A recursive depth first search.
///
/// Starting points are the nodes in the iterator `starts` (specify just one
/// start vertex *x* by using `Some(x)`).
///
/// The traversal emits discovery and finish events for each reachable vertex,
/// and edge classification of each reachable edge. `visitor` is called for each
/// event, see [`DfsEvent`][de] for possible values.
///
/// If the return value of the visitor is simply `()`, the visit runs until it
/// is finished. If the return value is a `Control<B>`, it can be used to
/// break the visit early, and the last control value is returned by the
/// function.
///
/// [de]: enum.DfsEvent.html
///
/// # Example
///
/// Find a path from vertex 0 to 5, and exit the visit as soon as we reach
/// the goal vertex.
///
/// ```
/// use petgraph::prelude::*;
/// use petgraph::graph::node_index as n;
/// use petgraph::visit::depth_first_search;
/// use petgraph::visit::{DfsEvent, Control};
///
/// let gr: Graph<(), ()> = Graph::from_edges(&[
///     (0, 1), (0, 2), (0, 3),
///     (1, 3),
///     (2, 3), (2, 4),
///     (4, 0), (4, 5),
/// ]);
///
/// // record each predecessor, mapping node → node
/// let mut predecessor = vec![NodeIndex::end(); gr.node_count()];
/// let start = n(0);
/// let goal = n(5);
/// depth_first_search(&gr, Some(start), |event| {
///     if let DfsEvent::TreeEdge(u, v) = event {
///         predecessor[v.index()] = u;
///         if v == goal {
///             return Control::Break(v);
///         }
///     }
///     Control::Continue
/// });
///
/// let mut next = goal;
/// let mut path = vec![next];
/// while next != start {
///     let pred = predecessor[next.index()];
///     path.push(pred);
///     next = pred;
/// }
/// path.reverse();
/// assert_eq!(&path, &[n(0), n(2), n(4), n(5)]);
/// ```
pub fn depth_first_search<G, I, F, C>(graph: G, starts: I, mut visitor: F) -> C
    where G: IntoNeighbors + Visitable,
          I: IntoIterator<Item=G::NodeId>,
          F: FnMut(DfsEvent<G::NodeId>) -> C,
          C: ControlFlow,
{
    let time = &mut Time(0);
    let discovered = &mut graph.visit_map();
    let finished = &mut graph.visit_map();

    for start in starts {
        try_control!(dfs_visitor(graph, start, &mut visitor, discovered, finished, time));
    }
    C::continuing()
}

fn dfs_visitor<G, F, C>(graph: G, u: G::NodeId, visitor: &mut F,
                     discovered: &mut G::Map, finished: &mut G::Map,
                     time: &mut Time) -> C
    where G: IntoNeighbors + Visitable,
          F: FnMut(DfsEvent<G::NodeId>) -> C,
          C: ControlFlow,
{
    if !discovered.visit(u) {
        return C::continuing();
    }
    try_control!(visitor(DfsEvent::Discover(u, time_post_inc(time))));
    for v in graph.neighbors(u) {
        if !discovered.is_visited(&v) {
            try_control!(visitor(DfsEvent::TreeEdge(u, v)));
            try_control!(dfs_visitor(graph, v, visitor, discovered, finished, time));
        } else if !finished.is_visited(&v) {
            try_control!(visitor(DfsEvent::BackEdge(u, v)));
        } else {
            try_control!(visitor(DfsEvent::CrossForwardEdge(u, v)));
        }
    }
    let first_finish = finished.visit(u);
    debug_assert!(first_finish);
    try_control!(visitor(DfsEvent::Finish(u, time_post_inc(time))));
    C::continuing()
}

fn time_post_inc(x: &mut Time) -> Time {
    let v = *x;
    x.0 += 1;
    v
}
<commit_msg>ControlFlow: No default implementation for should_break<commit_after>

use visit::IntoNeighbors;
use visit::{VisitMap, Visitable};

/// Strictly monotonically increasing event time for a depth first search.
#[derive(Copy, Clone, Debug, PartialEq, PartialOrd, Eq, Ord, Default)]
pub struct Time(pub usize);

/// A depth first search (DFS) visitor event.
#[derive(Copy, Clone, Debug)]
pub enum DfsEvent<N> {
    Discover(N, Time),
    /// An edge of the tree formed by the traversal.
    TreeEdge(N, N),
    /// An edge to an already visited node.
    BackEdge(N, N),
    /// A cross or forward edge.
    ///
    /// For an edge *(u, v)*, if the discover time of *v* is greater than *u*,
    /// then it is a forward edge, else a cross edge.
    CrossForwardEdge(N, N),
    Finish(N, Time),
}

/// Return if the expression is a break value.
macro_rules! try_control {
    ($e:expr) => {
        match $e {
            x => if x.should_break() {
                return x;
            }
        }
    }
}

/// Control flow for callbacks.
///
/// `Break` can carry a value.
#[derive(Copy, Clone, Debug)]
pub enum Control<B> {
    Continue,
    Break(B),
}

impl<B> Control<B> {
    pub fn breaking() -> Control<()> { Control::Break(()) }
    /// Get the value in `Control::Break(_)`, if present.
    pub fn break_value(self) -> Option<B> {
        match self {
            Control::Continue => None,
            Control::Break(b) => Some(b),
        }
    }
}

/// Control flow for callbacks.
///
/// The empty return value `()` is equivalent to continue.
pub trait ControlFlow {
    fn continuing() -> Self;
    fn should_break(&self) -> bool;
}

impl ControlFlow for () {
    fn continuing() { }
    #[inline]
    fn should_break(&self) -> bool { false }
}

impl<B> ControlFlow for Control<B> {
    fn continuing() -> Self { Control::Continue }
    fn should_break(&self) -> bool {
        if let Control::Break(_) = *self { true } else { false }
    }
}

impl<E> ControlFlow for Result<(), E> {
    fn continuing() -> Self { Ok(()) }
    fn should_break(&self) -> bool {
        if let Err(_) = *self { true } else { false }
    }
}

/// The default is `Continue`.
impl<B> Default for Control<B> {
    fn default() -> Self { Control::Continue }
}

/// A recursive depth first search.
///
/// Starting points are the nodes in the iterator `starts` (specify just one
/// start vertex *x* by using `Some(x)`).
///
/// The traversal emits discovery and finish events for each reachable vertex,
/// and edge classification of each reachable edge. `visitor` is called for each
/// event, see [`DfsEvent`][de] for possible values.
///
/// If the return value of the visitor is simply `()`, the visit runs until it
/// is finished. If the return value is a `Control<B>`, it can be used to
/// break the visit early, and the last control value is returned by the
/// function.
///
/// [de]: enum.DfsEvent.html
///
/// # Example
///
/// Find a path from vertex 0 to 5, and exit the visit as soon as we reach
/// the goal vertex.
///
/// ```
/// use petgraph::prelude::*;
/// use petgraph::graph::node_index as n;
/// use petgraph::visit::depth_first_search;
/// use petgraph::visit::{DfsEvent, Control};
///
/// let gr: Graph<(), ()> = Graph::from_edges(&[
///     (0, 1), (0, 2), (0, 3),
///     (1, 3),
///     (2, 3), (2, 4),
///     (4, 0), (4, 5),
/// ]);
///
/// // record each predecessor, mapping node → node
/// let mut predecessor = vec![NodeIndex::end(); gr.node_count()];
/// let start = n(0);
/// let goal = n(5);
/// depth_first_search(&gr, Some(start), |event| {
///     if let DfsEvent::TreeEdge(u, v) = event {
///         predecessor[v.index()] = u;
///         if v == goal {
///             return Control::Break(v);
///         }
///     }
///     Control::Continue
/// });
///
/// let mut next = goal;
/// let mut path = vec![next];
/// while next != start {
///     let pred = predecessor[next.index()];
///     path.push(pred);
///     next = pred;
/// }
/// path.reverse();
/// assert_eq!(&path, &[n(0), n(2), n(4), n(5)]);
/// ```
pub fn depth_first_search<G, I, F, C>(graph: G, starts: I, mut visitor: F) -> C
    where G: IntoNeighbors + Visitable,
          I: IntoIterator<Item=G::NodeId>,
          F: FnMut(DfsEvent<G::NodeId>) -> C,
          C: ControlFlow,
{
    let time = &mut Time(0);
    let discovered = &mut graph.visit_map();
    let finished = &mut graph.visit_map();

    for start in starts {
        try_control!(dfs_visitor(graph, start, &mut visitor, discovered, finished, time));
    }
    C::continuing()
}

fn dfs_visitor<G, F, C>(graph: G, u: G::NodeId, visitor: &mut F,
                     discovered: &mut G::Map, finished: &mut G::Map,
                     time: &mut Time) -> C
    where G: IntoNeighbors + Visitable,
          F: FnMut(DfsEvent<G::NodeId>) -> C,
          C: ControlFlow,
{
    if !discovered.visit(u) {
        return C::continuing();
    }
    try_control!(visitor(DfsEvent::Discover(u, time_post_inc(time))));
    for v in graph.neighbors(u) {
        if !discovered.is_visited(&v) {
            try_control!(visitor(DfsEvent::TreeEdge(u, v)));
            try_control!(dfs_visitor(graph, v, visitor, discovered, finished, time));
        } else if !finished.is_visited(&v) {
            try_control!(visitor(DfsEvent::BackEdge(u, v)));
        } else {
            try_control!(visitor(DfsEvent::CrossForwardEdge(u, v)));
        }
    }
    let first_finish = finished.visit(u);
    debug_assert!(first_finish);
    try_control!(visitor(DfsEvent::Finish(u, time_post_inc(time))));
    C::continuing()
}

fn time_post_inc(x: &mut Time) -> Time {
    let v = *x;
    x.0 += 1;
    v
}
<|endoftext|>


<commit_before>//! A simple crate for executing work on a thread pool, and getting back a
//! future.
//!
//! This crate provides a simple thread pool abstraction for running work
//! externally from the current thread that's running. An instance of `Future`
//! is handed back to represent that the work may be done later, and further
//! computations can be chained along with it as well.
//!
//! ```rust
//! extern crate futures;
//! extern crate futures_cpupool;
//!
//! use futures::Future;
//! use futures_cpupool::CpuPool;
//!
//! # fn long_running_future(a: u32) -> futures::future::BoxFuture<u32, ()> {
//! #     futures::future::result(Ok(a)).boxed()
//! # }
//! # fn main() {
//!
//! // Create a worker thread pool with four threads
//! let pool = CpuPool::new(4);
//!
//! // Execute some work on the thread pool, optionally closing over data.
//! let a = pool.spawn(long_running_future(2));
//! let b = pool.spawn(long_running_future(100));
//!
//! // Express some further computation once the work is completed on the thread
//! // pool.
//! let c = a.join(b).map(|(a, b)| a + b).wait().unwrap();
//!
//! // Print out the result
//! println!("{:?}", c);
//! # }
//! ```

#![deny(missing_docs)]

extern crate futures;
extern crate num_cpus;

use std::panic::{self, AssertUnwindSafe};
use std::sync::{Arc, Mutex};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::mpsc;
use std::thread;
use std::fmt;

use futures::{IntoFuture, Future, Poll, Async};
use futures::future::{lazy, Executor, ExecuteError};
use futures::sync::oneshot::{channel, Sender, Receiver};
use futures::executor::{self, Run, Executor as OldExecutor};

/// A thread pool intended to run CPU intensive work.
///
/// This thread pool will hand out futures representing the completed work
/// that happens on the thread pool itself, and the futures can then be later
/// composed with other work as part of an overall computation.
///
/// The worker threads associated with a thread pool are kept alive so long as
/// there is an open handle to the `CpuPool` or there is work running on them. Once
/// all work has been drained and all references have gone away the worker
/// threads will be shut down.
///
/// Currently `CpuPool` implements `Clone` which just clones a new reference to
/// the underlying thread pool.
///
/// **Note:** if you use CpuPool inside a library it's better accept a
/// `Builder` object for thread configuration rather than configuring just
/// pool size.  This not only future proof for other settings but also allows
/// user to attach monitoring tools to lifecycle hooks.
pub struct CpuPool {
    inner: Arc<Inner>,
}

/// Thread pool configuration object
///
/// Builder starts with a number of workers equal to the number
/// of CPUs on the host. But you can change it until you call `create()`.
pub struct Builder {
    pool_size: usize,
    name_prefix: Option<String>,
    after_start: Option<Arc<Fn() + Send + Sync>>,
    before_stop: Option<Arc<Fn() + Send + Sync>>,
}

struct MySender<F, T> {
    fut: F,
    tx: Option<Sender<T>>,
    keep_running_flag: Arc<AtomicBool>,
}

trait AssertSendSync: Send + Sync {}
impl AssertSendSync for CpuPool {}

struct Inner {
    tx: Mutex<mpsc::Sender<Message>>,
    rx: Mutex<mpsc::Receiver<Message>>,
    cnt: AtomicUsize,
    size: usize,
    after_start: Option<Arc<Fn() + Send + Sync>>,
    before_stop: Option<Arc<Fn() + Send + Sync>>,
}

impl fmt::Debug for CpuPool {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("CpuPool")
            .field("size", &self.inner.size)
            .finish()
    }
}

impl fmt::Debug for Builder {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Builder")
            .field("pool_size", &self.pool_size)
            .field("name_prefix", &self.name_prefix)
            .finish()
    }
}

/// The type of future returned from the `CpuPool::spawn` function, which
/// proxies the futures running on the thread pool.
///
/// This future will resolve in the same way as the underlying future, and it
/// will propagate panics.
#[must_use]
#[derive(Debug)]
pub struct CpuFuture<T, E> {
    inner: Receiver<thread::Result<Result<T, E>>>,
    keep_running_flag: Arc<AtomicBool>,
}

enum Message {
    Run(Run),
    Close,
}

impl CpuPool {
    /// Creates a new thread pool with `size` worker threads associated with it.
    ///
    /// The returned handle can use `execute` to run work on this thread pool,
    /// and clones can be made of it to get multiple references to the same
    /// thread pool.
    ///
    /// This is a shortcut for:
    ///
    /// ```rust
    /// # use futures_cpupool::{Builder, CpuPool};
    /// #
    /// # fn new(size: usize) -> CpuPool {
    /// Builder::new().pool_size(size).create()
    /// # }
    /// ```
    ///
    /// # Panics
    ///
    /// Panics if `size == 0`.
    pub fn new(size: usize) -> CpuPool {
        Builder::new().pool_size(size).create()
    }

    /// Creates a new thread pool with a number of workers equal to the number
    /// of CPUs on the host.
    ///
    /// This is a shortcut for:
    ///
    /// ```rust
    /// # use futures_cpupool::{Builder, CpuPool};
    /// #
    /// # fn new_num_cpus() -> CpuPool {
    /// Builder::new().create()
    /// # }
    /// ```
    pub fn new_num_cpus() -> CpuPool {
        Builder::new().create()
    }

    /// Spawns a future to run on this thread pool, returning a future
    /// representing the produced value.
    ///
    /// This function will execute the future `f` on the associated thread
    /// pool, and return a future representing the finished computation. The
    /// returned future serves as a proxy to the computation that `F` is
    /// running.
    ///
    /// To simply run an arbitrary closure on a thread pool and extract the
    /// result, you can use the `future::lazy` combinator to defer work to
    /// executing on the thread pool itself.
    ///
    /// Note that if the future `f` panics it will be caught by default and the
    /// returned future will propagate the panic. That is, panics will not tear
    /// down the thread pool and will be propagated to the returned future's
    /// `poll` method if queried.
    ///
    /// If the returned future is dropped then this `CpuPool` will attempt to
    /// cancel the computation, if possible. That is, if the computation is in
    /// the middle of working, it will be interrupted when possible.
    pub fn spawn<F>(&self, f: F) -> CpuFuture<F::Item, F::Error>
        where F: Future + Send + 'static,
              F::Item: Send + 'static,
              F::Error: Send + 'static,
    {
        let (tx, rx) = channel();
        let keep_running_flag = Arc::new(AtomicBool::new(false));
        // AssertUnwindSafe is used here becuase `Send + 'static` is basically
        // an alias for an implementation of the `UnwindSafe` trait but we can't
        // express that in the standard library right now.
        let sender = MySender {
            fut: AssertUnwindSafe(f).catch_unwind(),
            tx: Some(tx),
            keep_running_flag: keep_running_flag.clone(),
        };
        executor::spawn(sender).execute(self.inner.clone());
        CpuFuture { inner: rx , keep_running_flag: keep_running_flag.clone() }
    }

    /// Spawns a closure on this thread pool.
    ///
    /// This function is a convenience wrapper around the `spawn` function above
    /// for running a closure wrapped in `future::lazy`. It will spawn the
    /// function `f` provided onto the thread pool, and continue to run the
    /// future returned by `f` on the thread pool as well.
    ///
    /// The returned future will be a handle to the result produced by the
    /// future that `f` returns.
    pub fn spawn_fn<F, R>(&self, f: F) -> CpuFuture<R::Item, R::Error>
        where F: FnOnce() -> R + Send + 'static,
              R: IntoFuture + 'static,
              R::Future: Send + 'static,
              R::Item: Send + 'static,
              R::Error: Send + 'static,
    {
        self.spawn(lazy(f))
    }
}

impl<F> Executor<F> for CpuPool
    where F: Future<Item = (), Error = ()> + Send + 'static,
{
    fn execute(&self, future: F) -> Result<(), ExecuteError<F>> {
        executor::spawn(future).execute(self.inner.clone());
        Ok(())
    }
}

impl Inner {
    fn send(&self, msg: Message) {
        self.tx.lock().unwrap().send(msg).unwrap();
    }

    fn work(&self) {
        self.after_start.as_ref().map(|fun| fun());
        loop {
            let msg = self.rx.lock().unwrap().recv().unwrap();
            match msg {
                Message::Run(r) => r.run(),
                Message::Close => break,
            }
        }
        self.before_stop.as_ref().map(|fun| fun());
    }
}

impl Clone for CpuPool {
    fn clone(&self) -> CpuPool {
        self.inner.cnt.fetch_add(1, Ordering::Relaxed);
        CpuPool { inner: self.inner.clone() }
    }
}

impl Drop for CpuPool {
    fn drop(&mut self) {
        if self.inner.cnt.fetch_sub(1, Ordering::Relaxed) == 1 {
            for _ in 0..self.inner.size {
                self.inner.send(Message::Close);
            }
        }
    }
}

impl OldExecutor for Inner {
    fn execute(&self, run: Run) {
        self.send(Message::Run(run))
    }
}

impl<T, E> CpuFuture<T, E> {
    /// Drop this future without canceling the underlying future.
    ///
    /// When `CpuFuture` is dropped, `CpuPool` will try to abort the underlying
    /// future. This function can be used when user wants to drop but keep
    /// executing the underlying future.
    pub fn forget(self) {
        self.keep_running_flag.store(true, Ordering::SeqCst);
    }
}

impl<T: Send + 'static, E: Send + 'static> Future for CpuFuture<T, E> {
    type Item = T;
    type Error = E;

    fn poll(&mut self) -> Poll<T, E> {
        match self.inner.poll().expect("shouldn't be canceled") {
            Async::Ready(Ok(Ok(e))) => Ok(e.into()),
            Async::Ready(Ok(Err(e))) => Err(e),
            Async::Ready(Err(e)) => panic::resume_unwind(e),
            Async::NotReady => Ok(Async::NotReady),
        }
    }
}

impl<F: Future> Future for MySender<F, Result<F::Item, F::Error>> {
    type Item = ();
    type Error = ();

    fn poll(&mut self) -> Poll<(), ()> {
        if let Ok(Async::Ready(_)) = self.tx.as_mut().unwrap().poll_cancel() {
            if !self.keep_running_flag.load(Ordering::SeqCst) {
                // Cancelled, bail out
                return Ok(().into())
            }
        }

        let res = match self.fut.poll() {
            Ok(Async::Ready(e)) => Ok(e),
            Ok(Async::NotReady) => return Ok(Async::NotReady),
            Err(e) => Err(e),
        };

        // if the receiving end has gone away then that's ok, we just ignore the
        // send error here.
        drop(self.tx.take().unwrap().send(res));
        Ok(Async::Ready(()))
    }
}

impl Builder {
    /// Create a builder a number of workers equal to the number
    /// of CPUs on the host.
    pub fn new() -> Builder {
        Builder {
            pool_size: num_cpus::get(),
            name_prefix: None,
            after_start: None,
            before_stop: None,
        }
    }

    /// Set size of a future CpuPool
    ///
    /// The size of a thread pool is the number of worker threads spawned
    pub fn pool_size(&mut self, size: usize) -> &mut Self {
        self.pool_size = size;
        self
    }

    /// Set thread name prefix of a future CpuPool
    ///
    /// Thread name prefix is used for generating thread names. For example, if prefix is
    /// `my-pool-`, then threads in the pool will get names like `my-pool-1` etc.
    pub fn name_prefix<S: Into<String>>(&mut self, name_prefix: S) -> &mut Self {
        self.name_prefix = Some(name_prefix.into());
        self
    }

    /// Execute function `f` right after each thread is started but before
    /// running any jobs on it
    ///
    /// This is initially intended for bookkeeping and monitoring uses
    pub fn after_start<F>(&mut self, f: F) -> &mut Self
        where F: Fn() + Send + Sync + 'static
    {
        self.after_start = Some(Arc::new(f));
        self
    }

    /// Execute function `f` before each worker thread stops
    ///
    /// This is initially intended for bookkeeping and monitoring uses
    pub fn before_stop<F>(&mut self, f: F) -> &mut Self
        where F: Fn() + Send + Sync + 'static
    {
        self.before_stop = Some(Arc::new(f));
        self
    }

    /// Create CpuPool with configured parameters
    ///
    /// # Panics
    ///
    /// Panics if `pool_size == 0`.
    pub fn create(&mut self) -> CpuPool {
        let (tx, rx) = mpsc::channel();
        let pool = CpuPool {
            inner: Arc::new(Inner {
                tx: Mutex::new(tx),
                rx: Mutex::new(rx),
                cnt: AtomicUsize::new(1),
                size: self.pool_size,
                after_start: self.after_start.clone(),
                before_stop: self.before_stop.clone(),
            }),
        };
        assert!(self.pool_size > 0);

        for counter in 0..self.pool_size {
            let inner = pool.inner.clone();
            let mut thread_builder = thread::Builder::new();
            if let Some(ref name_prefix) = self.name_prefix {
                thread_builder = thread_builder.name(format!("{}{}", name_prefix, counter));
            }
            thread_builder.spawn(move || inner.work()).unwrap();
        }

        return pool
    }
}
<commit_msg>futures_cpupool: Deny missing `Debug` implementations<commit_after>//! A simple crate for executing work on a thread pool, and getting back a
//! future.
//!
//! This crate provides a simple thread pool abstraction for running work
//! externally from the current thread that's running. An instance of `Future`
//! is handed back to represent that the work may be done later, and further
//! computations can be chained along with it as well.
//!
//! ```rust
//! extern crate futures;
//! extern crate futures_cpupool;
//!
//! use futures::Future;
//! use futures_cpupool::CpuPool;
//!
//! # fn long_running_future(a: u32) -> futures::future::BoxFuture<u32, ()> {
//! #     futures::future::result(Ok(a)).boxed()
//! # }
//! # fn main() {
//!
//! // Create a worker thread pool with four threads
//! let pool = CpuPool::new(4);
//!
//! // Execute some work on the thread pool, optionally closing over data.
//! let a = pool.spawn(long_running_future(2));
//! let b = pool.spawn(long_running_future(100));
//!
//! // Express some further computation once the work is completed on the thread
//! // pool.
//! let c = a.join(b).map(|(a, b)| a + b).wait().unwrap();
//!
//! // Print out the result
//! println!("{:?}", c);
//! # }
//! ```

#![deny(missing_docs)]
#![deny(missing_debug_implementations)]

extern crate futures;
extern crate num_cpus;

use std::panic::{self, AssertUnwindSafe};
use std::sync::{Arc, Mutex};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::mpsc;
use std::thread;
use std::fmt;

use futures::{IntoFuture, Future, Poll, Async};
use futures::future::{lazy, Executor, ExecuteError};
use futures::sync::oneshot::{channel, Sender, Receiver};
use futures::executor::{self, Run, Executor as OldExecutor};

/// A thread pool intended to run CPU intensive work.
///
/// This thread pool will hand out futures representing the completed work
/// that happens on the thread pool itself, and the futures can then be later
/// composed with other work as part of an overall computation.
///
/// The worker threads associated with a thread pool are kept alive so long as
/// there is an open handle to the `CpuPool` or there is work running on them. Once
/// all work has been drained and all references have gone away the worker
/// threads will be shut down.
///
/// Currently `CpuPool` implements `Clone` which just clones a new reference to
/// the underlying thread pool.
///
/// **Note:** if you use CpuPool inside a library it's better accept a
/// `Builder` object for thread configuration rather than configuring just
/// pool size.  This not only future proof for other settings but also allows
/// user to attach monitoring tools to lifecycle hooks.
pub struct CpuPool {
    inner: Arc<Inner>,
}

/// Thread pool configuration object
///
/// Builder starts with a number of workers equal to the number
/// of CPUs on the host. But you can change it until you call `create()`.
pub struct Builder {
    pool_size: usize,
    name_prefix: Option<String>,
    after_start: Option<Arc<Fn() + Send + Sync>>,
    before_stop: Option<Arc<Fn() + Send + Sync>>,
}

struct MySender<F, T> {
    fut: F,
    tx: Option<Sender<T>>,
    keep_running_flag: Arc<AtomicBool>,
}

trait AssertSendSync: Send + Sync {}
impl AssertSendSync for CpuPool {}

struct Inner {
    tx: Mutex<mpsc::Sender<Message>>,
    rx: Mutex<mpsc::Receiver<Message>>,
    cnt: AtomicUsize,
    size: usize,
    after_start: Option<Arc<Fn() + Send + Sync>>,
    before_stop: Option<Arc<Fn() + Send + Sync>>,
}

impl fmt::Debug for CpuPool {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("CpuPool")
            .field("size", &self.inner.size)
            .finish()
    }
}

impl fmt::Debug for Builder {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Builder")
            .field("pool_size", &self.pool_size)
            .field("name_prefix", &self.name_prefix)
            .finish()
    }
}

/// The type of future returned from the `CpuPool::spawn` function, which
/// proxies the futures running on the thread pool.
///
/// This future will resolve in the same way as the underlying future, and it
/// will propagate panics.
#[must_use]
#[derive(Debug)]
pub struct CpuFuture<T, E> {
    inner: Receiver<thread::Result<Result<T, E>>>,
    keep_running_flag: Arc<AtomicBool>,
}

enum Message {
    Run(Run),
    Close,
}

impl CpuPool {
    /// Creates a new thread pool with `size` worker threads associated with it.
    ///
    /// The returned handle can use `execute` to run work on this thread pool,
    /// and clones can be made of it to get multiple references to the same
    /// thread pool.
    ///
    /// This is a shortcut for:
    ///
    /// ```rust
    /// # use futures_cpupool::{Builder, CpuPool};
    /// #
    /// # fn new(size: usize) -> CpuPool {
    /// Builder::new().pool_size(size).create()
    /// # }
    /// ```
    ///
    /// # Panics
    ///
    /// Panics if `size == 0`.
    pub fn new(size: usize) -> CpuPool {
        Builder::new().pool_size(size).create()
    }

    /// Creates a new thread pool with a number of workers equal to the number
    /// of CPUs on the host.
    ///
    /// This is a shortcut for:
    ///
    /// ```rust
    /// # use futures_cpupool::{Builder, CpuPool};
    /// #
    /// # fn new_num_cpus() -> CpuPool {
    /// Builder::new().create()
    /// # }
    /// ```
    pub fn new_num_cpus() -> CpuPool {
        Builder::new().create()
    }

    /// Spawns a future to run on this thread pool, returning a future
    /// representing the produced value.
    ///
    /// This function will execute the future `f` on the associated thread
    /// pool, and return a future representing the finished computation. The
    /// returned future serves as a proxy to the computation that `F` is
    /// running.
    ///
    /// To simply run an arbitrary closure on a thread pool and extract the
    /// result, you can use the `future::lazy` combinator to defer work to
    /// executing on the thread pool itself.
    ///
    /// Note that if the future `f` panics it will be caught by default and the
    /// returned future will propagate the panic. That is, panics will not tear
    /// down the thread pool and will be propagated to the returned future's
    /// `poll` method if queried.
    ///
    /// If the returned future is dropped then this `CpuPool` will attempt to
    /// cancel the computation, if possible. That is, if the computation is in
    /// the middle of working, it will be interrupted when possible.
    pub fn spawn<F>(&self, f: F) -> CpuFuture<F::Item, F::Error>
        where F: Future + Send + 'static,
              F::Item: Send + 'static,
              F::Error: Send + 'static,
    {
        let (tx, rx) = channel();
        let keep_running_flag = Arc::new(AtomicBool::new(false));
        // AssertUnwindSafe is used here becuase `Send + 'static` is basically
        // an alias for an implementation of the `UnwindSafe` trait but we can't
        // express that in the standard library right now.
        let sender = MySender {
            fut: AssertUnwindSafe(f).catch_unwind(),
            tx: Some(tx),
            keep_running_flag: keep_running_flag.clone(),
        };
        executor::spawn(sender).execute(self.inner.clone());
        CpuFuture { inner: rx , keep_running_flag: keep_running_flag.clone() }
    }

    /// Spawns a closure on this thread pool.
    ///
    /// This function is a convenience wrapper around the `spawn` function above
    /// for running a closure wrapped in `future::lazy`. It will spawn the
    /// function `f` provided onto the thread pool, and continue to run the
    /// future returned by `f` on the thread pool as well.
    ///
    /// The returned future will be a handle to the result produced by the
    /// future that `f` returns.
    pub fn spawn_fn<F, R>(&self, f: F) -> CpuFuture<R::Item, R::Error>
        where F: FnOnce() -> R + Send + 'static,
              R: IntoFuture + 'static,
              R::Future: Send + 'static,
              R::Item: Send + 'static,
              R::Error: Send + 'static,
    {
        self.spawn(lazy(f))
    }
}

impl<F> Executor<F> for CpuPool
    where F: Future<Item = (), Error = ()> + Send + 'static,
{
    fn execute(&self, future: F) -> Result<(), ExecuteError<F>> {
        executor::spawn(future).execute(self.inner.clone());
        Ok(())
    }
}

impl Inner {
    fn send(&self, msg: Message) {
        self.tx.lock().unwrap().send(msg).unwrap();
    }

    fn work(&self) {
        self.after_start.as_ref().map(|fun| fun());
        loop {
            let msg = self.rx.lock().unwrap().recv().unwrap();
            match msg {
                Message::Run(r) => r.run(),
                Message::Close => break,
            }
        }
        self.before_stop.as_ref().map(|fun| fun());
    }
}

impl Clone for CpuPool {
    fn clone(&self) -> CpuPool {
        self.inner.cnt.fetch_add(1, Ordering::Relaxed);
        CpuPool { inner: self.inner.clone() }
    }
}

impl Drop for CpuPool {
    fn drop(&mut self) {
        if self.inner.cnt.fetch_sub(1, Ordering::Relaxed) == 1 {
            for _ in 0..self.inner.size {
                self.inner.send(Message::Close);
            }
        }
    }
}

impl OldExecutor for Inner {
    fn execute(&self, run: Run) {
        self.send(Message::Run(run))
    }
}

impl<T, E> CpuFuture<T, E> {
    /// Drop this future without canceling the underlying future.
    ///
    /// When `CpuFuture` is dropped, `CpuPool` will try to abort the underlying
    /// future. This function can be used when user wants to drop but keep
    /// executing the underlying future.
    pub fn forget(self) {
        self.keep_running_flag.store(true, Ordering::SeqCst);
    }
}

impl<T: Send + 'static, E: Send + 'static> Future for CpuFuture<T, E> {
    type Item = T;
    type Error = E;

    fn poll(&mut self) -> Poll<T, E> {
        match self.inner.poll().expect("shouldn't be canceled") {
            Async::Ready(Ok(Ok(e))) => Ok(e.into()),
            Async::Ready(Ok(Err(e))) => Err(e),
            Async::Ready(Err(e)) => panic::resume_unwind(e),
            Async::NotReady => Ok(Async::NotReady),
        }
    }
}

impl<F: Future> Future for MySender<F, Result<F::Item, F::Error>> {
    type Item = ();
    type Error = ();

    fn poll(&mut self) -> Poll<(), ()> {
        if let Ok(Async::Ready(_)) = self.tx.as_mut().unwrap().poll_cancel() {
            if !self.keep_running_flag.load(Ordering::SeqCst) {
                // Cancelled, bail out
                return Ok(().into())
            }
        }

        let res = match self.fut.poll() {
            Ok(Async::Ready(e)) => Ok(e),
            Ok(Async::NotReady) => return Ok(Async::NotReady),
            Err(e) => Err(e),
        };

        // if the receiving end has gone away then that's ok, we just ignore the
        // send error here.
        drop(self.tx.take().unwrap().send(res));
        Ok(Async::Ready(()))
    }
}

impl Builder {
    /// Create a builder a number of workers equal to the number
    /// of CPUs on the host.
    pub fn new() -> Builder {
        Builder {
            pool_size: num_cpus::get(),
            name_prefix: None,
            after_start: None,
            before_stop: None,
        }
    }

    /// Set size of a future CpuPool
    ///
    /// The size of a thread pool is the number of worker threads spawned
    pub fn pool_size(&mut self, size: usize) -> &mut Self {
        self.pool_size = size;
        self
    }

    /// Set thread name prefix of a future CpuPool
    ///
    /// Thread name prefix is used for generating thread names. For example, if prefix is
    /// `my-pool-`, then threads in the pool will get names like `my-pool-1` etc.
    pub fn name_prefix<S: Into<String>>(&mut self, name_prefix: S) -> &mut Self {
        self.name_prefix = Some(name_prefix.into());
        self
    }

    /// Execute function `f` right after each thread is started but before
    /// running any jobs on it
    ///
    /// This is initially intended for bookkeeping and monitoring uses
    pub fn after_start<F>(&mut self, f: F) -> &mut Self
        where F: Fn() + Send + Sync + 'static
    {
        self.after_start = Some(Arc::new(f));
        self
    }

    /// Execute function `f` before each worker thread stops
    ///
    /// This is initially intended for bookkeeping and monitoring uses
    pub fn before_stop<F>(&mut self, f: F) -> &mut Self
        where F: Fn() + Send + Sync + 'static
    {
        self.before_stop = Some(Arc::new(f));
        self
    }

    /// Create CpuPool with configured parameters
    ///
    /// # Panics
    ///
    /// Panics if `pool_size == 0`.
    pub fn create(&mut self) -> CpuPool {
        let (tx, rx) = mpsc::channel();
        let pool = CpuPool {
            inner: Arc::new(Inner {
                tx: Mutex::new(tx),
                rx: Mutex::new(rx),
                cnt: AtomicUsize::new(1),
                size: self.pool_size,
                after_start: self.after_start.clone(),
                before_stop: self.before_stop.clone(),
            }),
        };
        assert!(self.pool_size > 0);

        for counter in 0..self.pool_size {
            let inner = pool.inner.clone();
            let mut thread_builder = thread::Builder::new();
            if let Some(ref name_prefix) = self.name_prefix {
                thread_builder = thread_builder.name(format!("{}{}", name_prefix, counter));
            }
            thread_builder.spawn(move || inner.work()).unwrap();
        }

        return pool
    }
}
<|endoftext|>


<commit_before>#[macro_use] extern crate log;
extern crate lapin_futures as lapin;
extern crate futures;
extern crate tokio_core;
extern crate env_logger;

use futures::future::Future;
use futures::Stream;
use tokio_core::reactor::Core;
use tokio_core::net::TcpStream;
use lapin::client::ConnectionOptions;
use lapin::channel::{BasicConsumeOptions,BasicGetOptions,BasicPublishOptions,BasicProperties,QueueDeclareOptions};

fn main() {
  env_logger::init().unwrap();
  let mut core = Core::new().unwrap();

  let handle = core.handle();
  let addr = "127.0.0.1:5672".parse().unwrap();

  core.run(
    TcpStream::connect(&addr, &handle).and_then(|stream| {
      lapin::client::Client::connect(stream, &ConnectionOptions::default())
    }).and_then(|client| {

      client.create_channel().and_then(|channel| {
        let id = channel.id;
        info!("created channel with id: {}", id);

        channel.queue_declare("hello", &QueueDeclareOptions::default()).and_then(move |_| {
          info!("channel {} declared queue {}", id, "hello");

          channel.basic_publish("hello", b"hello from tokio", &BasicPublishOptions::default(), BasicProperties::default().with_user_id("guest".to_string()).with_reply_to("foobar".to_string()))
        })
      }).and_then(move |_| {
        client.create_channel()
      }).and_then(|channel| {
        let id = channel.id;
        info!("created channel with id: {}", id);

        let c = channel.clone();
        channel.queue_declare("hello", &QueueDeclareOptions::default()).and_then(move |_| {
          info!("channel {} declared queue {}", id, "hello");

          channel.basic_consume("hello", "my_consumer", &BasicConsumeOptions::default())
        }).and_then(|stream| {
          info!("got consumer stream");

          stream.for_each(move |message| {
            debug!("got message: {:?}", message);
            info!("decoded message: {:?}", std::str::from_utf8(&message.data).unwrap());
            c.basic_ack(message.delivery_tag);
            Ok(())
          })
        })
      })
    })
  ).unwrap();
}
<commit_msg>use a confirm channel in the example<commit_after>#[macro_use] extern crate log;
extern crate lapin_futures as lapin;
extern crate futures;
extern crate tokio_core;
extern crate env_logger;

use futures::future::Future;
use futures::Stream;
use tokio_core::reactor::Core;
use tokio_core::net::TcpStream;
use lapin::client::ConnectionOptions;
use lapin::channel::{BasicConsumeOptions,BasicGetOptions,BasicPublishOptions,BasicProperties,QueueDeclareOptions};

fn main() {
  env_logger::init().unwrap();
  let mut core = Core::new().unwrap();

  let handle = core.handle();
  let addr = "127.0.0.1:5672".parse().unwrap();

  core.run(
    TcpStream::connect(&addr, &handle).and_then(|stream| {
      lapin::client::Client::connect(stream, &ConnectionOptions::default())
    }).and_then(|client| {

      client.create_confirm_channel().and_then(|channel| {
        let id = channel.id;
        info!("created channel with id: {}", id);

        channel.queue_declare("hello", &QueueDeclareOptions::default()).and_then(move |_| {
          info!("channel {} declared queue {}", id, "hello");

          channel.basic_publish(
            "hello",
            b"hello from tokio",
            &BasicPublishOptions::default(),
            BasicProperties::default().with_user_id("guest".to_string()).with_reply_to("foobar".to_string())
          ).map(|confirmation| {
            info!("publish got confirmation: {:?}", confirmation)
          })
        })
      }).and_then(move |_| {
        client.create_channel()
      }).and_then(|channel| {
        let id = channel.id;
        info!("created channel with id: {}", id);

        let c = channel.clone();
        channel.queue_declare("hello", &QueueDeclareOptions::default()).and_then(move |_| {
          info!("channel {} declared queue {}", id, "hello");

          channel.basic_consume("hello", "my_consumer", &BasicConsumeOptions::default())
        }).and_then(|stream| {
          info!("got consumer stream");

          stream.for_each(move |message| {
            debug!("got message: {:?}", message);
            info!("decoded message: {:?}", std::str::from_utf8(&message.data).unwrap());
            c.basic_ack(message.delivery_tag);
            Ok(())
          })
        })
      })
    })
  ).unwrap();
}
<|endoftext|>