<commit_before>pub use self::parse_mode::ParseMode;
pub use self::chat_action::ChatAction;

mod parse_mode;
mod chat_action;
mod file;

use reqwest::Client;
use serde_json;
use serde_json::Value;

use bot::chat_action::get_chat_action;
use bot::parse_mode::get_parse_mode;
use error::{Result, check_for_error};
use error::Error::JsonNotFound;
use marker::ReplyMarkup;
use objects::{Update, Message, Contact, InlineKeyboardMarkup, User};

/// A `Bot` which will do all the API calls.
///
/// The `Bot` will be given access to in a `Command` with which you can do all
/// the API interactions in your `Command`s.
#[derive(Debug)]
pub struct Bot {
    pub id: i64,
    pub first_name: String,
    pub last_name: Option<String>,
    pub username: String,
    pub bot_url: String,
    client: Client,
}

impl Bot {
    /// Constructs a new `Bot`.
    pub fn new(bot_url: String) -> Result<Self> {
        debug!("Going to construct a new Bot...");
        let client = Client::new()?;
        let me = Bot::get_me(&client, &bot_url)?;
        let id = me.id;
        let first_name = me.first_name;
        let last_name = me.last_name;
        let username = me.username.expect("Cannot find username of the bot");

        Ok(Bot {
            id: id,
            first_name: first_name,
            last_name: last_name,
            username: username,
            client: client,
            bot_url: bot_url,
        })
    }

    /// API call which gets the information about your bot.
    pub fn get_me(client: &Client, bot_url: &str) -> Result<User> {
        debug!("Calling get_me...");
        let path = ["getMe"];
        let url = ::construct_api_url(bot_url, &path);
        let mut data = client.get(&url)?.send()?;
        let rjson: Value = check_for_error(data.json()?)?;
        let user_json = rjson.get("result").ok_or(JsonNotFound)?;
        let user: User = serde_json::from_value(user_json.clone())?;
        Ok(user)
    }

    /// API call which will get called to get the updates for your bot.
    pub fn get_updates(
        &self,
        offset: i32,
        limit: Option<i32>,
        timeout: Option<i32>,
        network_delay: Option<f32>,
    ) -> Result<Option<Vec<Update>>> {
        debug!("Calling get_updates...");
        let limit = limit.unwrap_or(100);
        let timeout = timeout.unwrap_or(0);
        let network_delay = network_delay.unwrap_or(0.0);
        let path = ["getUpdates"];
        let path_url = ::construct_api_url(&self.bot_url, &path);
        let url = format!(
            "{}?offset={}&limit={}&timeout={}&network_delay={}",
            path_url,
            offset,
            limit,
            timeout,
            network_delay
        );
        let mut data = self.client.get(&url)?.send()?;
        let rjson: Value = check_for_error(data.json()?)?;
        let updates_json = rjson.get("result");

        if let Some(result) = updates_json {
            let updates: Vec<Update> = serde_json::from_value(result.clone())?;
            Ok(Some(updates))
        } else {
            Ok(None)
        }
    }

    /// API call which will send a message to a chat which your bot participates in.
    pub fn send_message<M: ReplyMarkup>(
        &self,
        chat_id: &i64,
        text: &str,
        parse_mode: Option<&ParseMode>,
        disable_web_page_preview: Option<&bool>,
        disable_notification: Option<&bool>,
        reply_to_message_id: Option<&i64>,
        reply_markup: Option<M>,
    ) -> Result<Message> {
        debug!("Calling send_message...");
        let chat_id: &str = &chat_id.to_string();
        let parse_mode = &get_parse_mode(parse_mode.unwrap_or(&ParseMode::Text));
        let disable_web_page_preview: &str =
            &disable_web_page_preview.unwrap_or(&false).to_string();
        let disable_notification: &str = &disable_notification.unwrap_or(&false).to_string();
        let reply_to_message_id: &str = &reply_to_message_id.map(|i| i.to_string()).unwrap_or(
            "None"
                .to_string(),
        );
        let reply_markup = &Box::new(reply_markup)
            .map(|r| serde_json::to_string(&r).unwrap_or("".to_string()))
            .unwrap_or("".to_string());

        let path = ["sendMessage"];
        let params = [
            ("chat_id", chat_id),
            ("text", text),
            ("parse_mode", parse_mode),
            ("disable_web_page_preview", disable_web_page_preview),
            ("disable_notification", disable_notification),
            ("reply_to_message_id", reply_to_message_id),
            ("reply_markup", reply_markup),
        ];
        self.post_message(&path, &params)
    }

    /// API call which will reply to a message directed to your bot.
    pub fn reply_to_message(&self, update: &Update, text: &str) -> Result<Message> {
        debug!("Calling reply_to_message...");
        let message = update.clone().message.unwrap();
        let message_id = message.message_id;
        let chat_id = message.chat.id;
        self.send_message(
            &chat_id,
            text,
            None,
            None,
            None,
            Some(&message_id),
            ::NO_MARKUP,
        )
    }

    /// API call which will forward a message.
    pub fn forward_message(
        &self,
        update: &Update,
        chat_id: &i64,
        disable_notification: Option<&bool>,
    ) -> Result<Message> {
        debug!("Calling forward_message...");
        let message = update.clone().message.unwrap();
        let chat_id: &str = &chat_id.to_string();
        let from_chat_id: &str = &message.chat.id.to_string();
        let message_id: &str = &message.message_id.to_string();
        let disable_notification: &str = &disable_notification.unwrap_or(&false).to_string();
        let path = ["forwardMessage"];
        let params = [
            ("chat_id", chat_id),
            ("from_chat_id", from_chat_id),
            ("disable_notification", disable_notification),
            ("message_id", message_id),
        ];
        self.post_message(&path, &params)
    }

    /// API call which will show the given chat action to the users.
    pub fn send_chat_action(&self, chat_id: &i64, action: &ChatAction) -> Result<bool> {
        debug!("Calling send_chat_action...");
        let chat_id: &str = &chat_id.to_string();
        let action = &get_chat_action(action);
        let path = ["sendChatAction"];
        let params = [("chat_id", chat_id), ("action", action)];
        let url = ::construct_api_url(&self.bot_url, &path);
        let mut data = self.client.post(&url)?.form(&params)?.send()?;
        let rjson: Value = check_for_error(data.json()?)?;
        let result_json = rjson.get("result").ok_or(JsonNotFound)?;
        let chat_action_succeeded: bool = serde_json::from_value(result_json.clone())?;
        Ok(chat_action_succeeded)
    }

    /// API call which will send the given contact.
    pub fn send_contact(
        &self,
        chat_id: &i64,
        contact: &Contact,
        disable_notification: Option<&bool>,
        reply_to_message_id: Option<&i64>,
        reply_markup: Option<&InlineKeyboardMarkup>,
    ) -> Result<Message> {
        debug!("Calling send_contact...");
        let chat_id: &str = &chat_id.to_string();
        let phone_number = &contact.phone_number;
        let first_name = &contact.first_name;
        let last_name = &contact.clone().last_name.unwrap();
        let disable_notification: &str = &disable_notification.unwrap_or(&false).to_string();
        let reply_to_message_id: &str = &reply_to_message_id.map(|i| i.to_string()).unwrap_or(
            "None"
                .to_string(),
        );
        let reply_markup = &reply_markup
            .and_then(|r| serde_json::to_string(r).ok())
            .unwrap_or("".to_string());
        let path = ["sendContact"];
        let params = [
            ("chat_id", chat_id),
            ("phone_number", phone_number),
            ("first_name", first_name),
            ("last_name", last_name),
            ("disable_notification", disable_notification),
            ("reply_to_message_id", reply_to_message_id),
            ("reply_markup", reply_markup),
        ];
        self.post_message(&path, &params)
    }

    /// The actual networking done for sending messages.
    fn post_message(&self, path: &[&str], params: &[(&str, &str)]) -> Result<Message> {
        debug!("Posting message...");
        let url = ::construct_api_url(&self.bot_url, path);
        let mut data = self.client.post(&url)?.form(&params)?.send()?;
        let rjson: Value = check_for_error(data.json()?)?;
        let message_json = rjson.get("result").ok_or(JsonNotFound)?;
        let message: Message = serde_json::from_value(message_json.clone())?;
        Ok(message)
    }
}
<commit_msg>Making post_message generic<commit_after>pub use self::parse_mode::ParseMode;
pub use self::chat_action::ChatAction;

mod parse_mode;
mod chat_action;
mod file;

use reqwest::Client;
use serde::de::DeserializeOwned;
use serde_json;
use serde_json::Value;

use bot::chat_action::get_chat_action;
use bot::parse_mode::get_parse_mode;
use error::{Result, check_for_error};
use error::Error::JsonNotFound;
use marker::ReplyMarkup;
use objects::{Update, Message, Contact, InlineKeyboardMarkup, User, UserProfilePhotos};

/// A `Bot` which will do all the API calls.
///
/// The `Bot` will be given access to in a `Command` with which you can do all
/// the API interactions in your `Command`s.
#[derive(Debug)]
pub struct Bot {
    pub id: i64,
    pub first_name: String,
    pub last_name: Option<String>,
    pub username: String,
    pub bot_url: String,
    client: Client,
}

impl Bot {
    /// Constructs a new `Bot`.
    pub fn new(bot_url: String) -> Result<Self> {
        debug!("Going to construct a new Bot...");
        let client = Client::new()?;
        let me = Bot::get_me(&client, &bot_url)?;
        let id = me.id;
        let first_name = me.first_name;
        let last_name = me.last_name;
        let username = me.username.expect("Cannot find username of the bot");

        Ok(Bot {
            id: id,
            first_name: first_name,
            last_name: last_name,
            username: username,
            client: client,
            bot_url: bot_url,
        })
    }

    /// API call which gets the information about your bot.
    pub fn get_me(client: &Client, bot_url: &str) -> Result<User> {
        debug!("Calling get_me...");
        let path = ["getMe"];
        let url = ::construct_api_url(bot_url, &path);
        let mut data = client.get(&url)?.send()?;
        let rjson: Value = check_for_error(data.json()?)?;
        let user_json = rjson.get("result").ok_or(JsonNotFound)?;
        let user: User = serde_json::from_value(user_json.clone())?;
        Ok(user)
    }

    /// API call which will get called to get the updates for your bot.
    pub fn get_updates(
        &self,
        offset: i32,
        limit: Option<i32>,
        timeout: Option<i32>,
        network_delay: Option<f32>,
    ) -> Result<Option<Vec<Update>>> {
        debug!("Calling get_updates...");
        let limit = limit.unwrap_or(100);
        let timeout = timeout.unwrap_or(0);
        let network_delay = network_delay.unwrap_or(0.0);
        let path = ["getUpdates"];
        let path_url = ::construct_api_url(&self.bot_url, &path);
        let url = format!(
            "{}?offset={}&limit={}&timeout={}&network_delay={}",
            path_url,
            offset,
            limit,
            timeout,
            network_delay
        );
        let mut data = self.client.get(&url)?.send()?;
        let rjson: Value = check_for_error(data.json()?)?;
        let updates_json = rjson.get("result");

        if let Some(result) = updates_json {
            let updates: Vec<Update> = serde_json::from_value(result.clone())?;
            Ok(Some(updates))
        } else {
            Ok(None)
        }
    }

    /// API call which will send a message to a chat which your bot participates in.
    pub fn send_message<M: ReplyMarkup>(
        &self,
        chat_id: &i64,
        text: &str,
        parse_mode: Option<&ParseMode>,
        disable_web_page_preview: Option<&bool>,
        disable_notification: Option<&bool>,
        reply_to_message_id: Option<&i64>,
        reply_markup: Option<M>,
    ) -> Result<Message> {
        debug!("Calling send_message...");
        let chat_id: &str = &chat_id.to_string();
        let parse_mode = &get_parse_mode(parse_mode.unwrap_or(&ParseMode::Text));
        let disable_web_page_preview: &str =
            &disable_web_page_preview.unwrap_or(&false).to_string();
        let disable_notification: &str = &disable_notification.unwrap_or(&false).to_string();
        let reply_to_message_id: &str = &reply_to_message_id.map(|i| i.to_string()).unwrap_or(
            "None"
                .to_string(),
        );
        let reply_markup = &Box::new(reply_markup)
            .map(|r| serde_json::to_string(&r).unwrap_or("".to_string()))
            .unwrap_or("".to_string());

        let path = ["sendMessage"];
        let params = [
            ("chat_id", chat_id),
            ("text", text),
            ("parse_mode", parse_mode),
            ("disable_web_page_preview", disable_web_page_preview),
            ("disable_notification", disable_notification),
            ("reply_to_message_id", reply_to_message_id),
            ("reply_markup", reply_markup),
        ];
        self.post_message(&path, &params)
    }

    /// API call which will reply to a message directed to your bot.
    pub fn reply_to_message(&self, update: &Update, text: &str) -> Result<Message> {
        debug!("Calling reply_to_message...");
        let message = update.clone().message.unwrap();
        let message_id = message.message_id;
        let chat_id = message.chat.id;
        self.send_message(
            &chat_id,
            text,
            None,
            None,
            None,
            Some(&message_id),
            ::NO_MARKUP,
        )
    }

    /// API call which will forward a message.
    pub fn forward_message(
        &self,
        update: &Update,
        chat_id: &i64,
        disable_notification: Option<&bool>,
    ) -> Result<Message> {
        debug!("Calling forward_message...");
        let message = update.clone().message.unwrap();
        let chat_id: &str = &chat_id.to_string();
        let from_chat_id: &str = &message.chat.id.to_string();
        let message_id: &str = &message.message_id.to_string();
        let disable_notification: &str = &disable_notification.unwrap_or(&false).to_string();
        let path = ["forwardMessage"];
        let params = [
            ("chat_id", chat_id),
            ("from_chat_id", from_chat_id),
            ("disable_notification", disable_notification),
            ("message_id", message_id),
        ];
        self.post_message(&path, &params)
    }

    /// API call which will show the given chat action to the users.
    pub fn send_chat_action(&self, chat_id: &i64, action: &ChatAction) -> Result<bool> {
        debug!("Calling send_chat_action...");
        let chat_id: &str = &chat_id.to_string();
        let action = &get_chat_action(action);
        let path = ["sendChatAction"];
        let params = [("chat_id", chat_id), ("action", action)];
        self.post_message(&path, &params)
    }

    /// API call which will send the given contact.
    pub fn send_contact(
        &self,
        chat_id: &i64,
        contact: &Contact,
        disable_notification: Option<&bool>,
        reply_to_message_id: Option<&i64>,
        reply_markup: Option<&InlineKeyboardMarkup>,
    ) -> Result<Message> {
        debug!("Calling send_contact...");
        let chat_id: &str = &chat_id.to_string();
        let phone_number = &contact.phone_number;
        let first_name = &contact.first_name;
        let last_name = &contact.clone().last_name.unwrap();
        let disable_notification: &str = &disable_notification.unwrap_or(&false).to_string();
        let reply_to_message_id: &str = &reply_to_message_id.map(|i| i.to_string()).unwrap_or(
            "None"
                .to_string(),
        );
        let reply_markup = &reply_markup
            .and_then(|r| serde_json::to_string(r).ok())
            .unwrap_or("".to_string());
        let path = ["sendContact"];
        let params = [
            ("chat_id", chat_id),
            ("phone_number", phone_number),
            ("first_name", first_name),
            ("last_name", last_name),
            ("disable_notification", disable_notification),
            ("reply_to_message_id", reply_to_message_id),
            ("reply_markup", reply_markup),
        ];
        self.post_message(&path, &params)
    }

    pub fn get_user_profile_photos(
        &self,
        user_id: &i64,
        offset: Option<i64>,
        limit: Option<i64>,
    ) -> Result<UserProfilePhotos> {
        let user_id: &str = &user_id.to_string();
        let offset: &str = &offset.map(|i| i.to_string()).unwrap_or("None".to_string());
        let limit: &str = &limit.map(|i| i.to_string()).unwrap_or("None".to_string());
        let path = ["getUserProfilePhotos"];
        let params = [("user_id", user_id), ("offset", offset), ("limit", limit)];
        self.post_message(&path, &params)
    }

    /// The actual networking done for sending messages.
    fn post_message<T: DeserializeOwned>(
        &self,
        path: &[&str],
        params: &[(&str, &str)],
    ) -> Result<T> {
        debug!("Posting message...");
        let url = ::construct_api_url(&self.bot_url, path);
        let mut data = self.client.post(&url)?.form(&params)?.send()?;
        let rjson: Value = check_for_error(data.json()?)?;
        let object_json = rjson.get("result").ok_or(JsonNotFound)?;
        let object = serde_json::from_value::<T>(object_json.clone())?;
        Ok(object)
    }
}
<|endoftext|>


<commit_before>// * This file is part of the uutils coreutils package.
// *
// * (c) 2015 Wiktor Kuropatwa <wiktor.kuropatwa@gmail.com>
// * (c) 2020 nicoo            <nicoo@debian.org>
// *
// * For the full copyright and license information, please view the LICENSE file
// * that was distributed with this source code.

use num_traits::{
    identities::{One, Zero},
    int::PrimInt,
    ops::wrapping::{WrappingMul, WrappingNeg, WrappingSub},
};
use std::fmt::{Debug, Display};
use std::mem::swap;

// This is incorrectly reported as dead code,
//  presumably when included in build.rs.
#[allow(dead_code)]
pub(crate) fn gcd(mut a: u64, mut b: u64) -> u64 {
    while b > 0 {
        a %= b;
        swap(&mut a, &mut b);
    }
    a
}

pub(crate) trait Arithmetic: Copy + Sized {
    // The type of integers mod m, in some opaque representation
    type ModInt: Copy + Sized + Eq;

    fn new(m: u64) -> Self;
    fn modulus(&self) -> u64;
    fn from_u64(&self, n: u64) -> Self::ModInt;
    fn to_u64(&self, n: Self::ModInt) -> u64;
    fn add(&self, a: Self::ModInt, b: Self::ModInt) -> Self::ModInt;
    fn mul(&self, a: Self::ModInt, b: Self::ModInt) -> Self::ModInt;

    fn pow(&self, mut a: Self::ModInt, mut b: u64) -> Self::ModInt {
        let (_a, _b) = (a, b);
        let mut result = self.one();
        while b > 0 {
            if b & 1 != 0 {
                result = self.mul(result, a);
            }
            a = self.mul(a, a);
            b >>= 1;
        }

        // Check that r (reduced back to the usual representation) equals
        //  a^b % n, unless the latter computation overflows
        // Temporarily commented-out, as there u64::checked_pow is not available
        //  on the minimum supported Rust version, nor is an appropriate method
        //  for compiling the check conditionally.
        //debug_assert!(self
        //    .to_u64(_a)
        //    .checked_pow(_b as u32)
        //    .map(|r| r % self.modulus() == self.to_u64(result))
        //    .unwrap_or(true));

        result
    }

    fn one(&self) -> Self::ModInt {
        self.from_u64(1)
    }
    fn minus_one(&self) -> Self::ModInt {
        self.from_u64(self.modulus() - 1)
    }
    fn zero(&self) -> Self::ModInt {
        self.from_u64(0)
    }
}

#[derive(Clone, Copy, Debug)]
pub(crate) struct Montgomery<T: DoubleInt> {
    a: T,
    n: T,
}

impl<T: DoubleInt> Montgomery<T> {
    /// computes x/R mod n efficiently
    fn reduce(&self, x: T::Double) -> T {
        let t_bits = T::zero().count_zeros() as usize;

        debug_assert!(x < (self.n.as_double()) << t_bits);
        // TODO: optimiiiiiiise
        let Montgomery { a, n } = self;
        let m = T::from_double(x).wrapping_mul(a);
        let nm = (n.as_double()) * (m.as_double());
        let (xnm, overflow) = x.overflowing_add_(nm); // x + n*m
        debug_assert_eq!(
            xnm % (T::Double::one() << T::zero().count_zeros() as usize),
            T::Double::zero()
        );

        // (x + n*m) / R
        // in case of overflow, this is (2¹²⁸ + xnm)/2⁶⁴ - n = xnm/2⁶⁴ + (2⁶⁴ - n)
        let y = T::from_double(xnm >> t_bits)
            + if !overflow {
                T::zero()
            } else {
                n.wrapping_neg()
            };

        if y >= *n {
            y - *n
        } else {
            y
        }
    }
}

impl<T: DoubleInt> Arithmetic for Montgomery<T> {
    // Montgomery transform, R=2⁶⁴
    // Provides fast arithmetic mod n (n odd, u64)
    type ModInt = T;

    fn new(n: u64) -> Self {
        debug_assert!(T::zero().count_zeros() >= 64 || n < (1 << T::zero().count_zeros() as usize));
        let n = T::from_u64(n);
        let a = modular_inverse(n).wrapping_neg();
        debug_assert_eq!(n.wrapping_mul(&a), T::one().wrapping_neg());
        Montgomery { a, n }
    }

    fn modulus(&self) -> u64 {
        self.n.as_u64()
    }

    fn from_u64(&self, x: u64) -> Self::ModInt {
        // TODO: optimise!
        debug_assert!(x < self.n.as_u64());
        let r = T::from_double(
            ((T::Double::from_u64(x)) << T::zero().count_zeros() as usize) % self.n.as_double(),
        );
        debug_assert_eq!(x, self.to_u64(r));
        r
    }

    fn to_u64(&self, n: Self::ModInt) -> u64 {
        self.reduce(n.as_double()).as_u64()
    }

    fn add(&self, a: Self::ModInt, b: Self::ModInt) -> Self::ModInt {
        let (r, overflow) = a.overflowing_add_(b);

        // In case of overflow, a+b = 2⁶⁴ + r = (2⁶⁴ - n) + r (working mod n)
        let r = if !overflow {
            r
        } else {
            r + self.n.wrapping_neg()
        };

        // Normalise to [0; n[
        let r = if r < self.n { r } else { r - self.n };

        // Check that r (reduced back to the usual representation) equals
        // a+b % n
        #[cfg(debug_assertions)]
        {
            let a_r = self.to_u64(a) as u128;
            let b_r = self.to_u64(b) as u128;
            let r_r = self.to_u64(r);
            let r_2 = ((a_r + b_r) % self.n.as_u128()) as u64;
            debug_assert_eq!(
                r_r, r_2,
                "[{}] = {} ≠ {} = {} + {} = [{}] + [{}] mod {}; a = {}",
                r, r_r, r_2, a_r, b_r, a, b, self.n, self.a
            );
        }
        r
    }

    fn mul(&self, a: Self::ModInt, b: Self::ModInt) -> Self::ModInt {
        let r = self.reduce(a.as_double() * b.as_double());

        // Check that r (reduced back to the usual representation) equals
        // a*b % n
        #[cfg(debug_assertions)]
        {
            let a_r = self.to_u64(a) as u128;
            let b_r = self.to_u64(b) as u128;
            let r_r = self.to_u64(r);
            let r_2: u64 = ((a_r * b_r) % self.n.as_u128()) as u64;
            debug_assert_eq!(
                r_r, r_2,
                "[{}] = {} ≠ {} = {} * {} = [{}] * [{}] mod {}; a = {}",
                r, r_r, r_2, a_r, b_r, a, b, self.n, self.a
            );
        }
        r
    }
}

// NOTE: Trait can be removed once num-traits adds a similar one;
//       see https://github.com/rust-num/num-traits/issues/168
pub(crate) trait OverflowingAdd: Sized {
    fn overflowing_add_(self, n: Self) -> (Self, bool);
}
macro_rules! overflowing {
    ($x:ty) => {
        impl OverflowingAdd for $x {
            fn overflowing_add_(self, n: Self) -> (Self, bool) {
                self.overflowing_add(n)
            }
        }
    };
}
overflowing!(u32);
overflowing!(u64);
overflowing!(u128);

pub(crate) trait Int:
    Display + Debug + PrimInt + OverflowingAdd + WrappingNeg + WrappingSub + WrappingMul
{
    fn as_u64(&self) -> u64;
    fn from_u64(n: u64) -> Self;
    #[cfg(debug_assertions)]
    fn as_u128(&self) -> u128;
    #[cfg(debug_assertions)]
    fn from_u128(n: u64) -> Self;
}

pub(crate) trait DoubleInt: Int {
    type Double: Int;

    fn as_double(self) -> Self::Double;
    fn from_double(n: Self::Double) -> Self;
}

macro_rules! int {
    ( $x:ty ) => {
        impl Int for $x {
            fn as_u64(&self) -> u64 {
                *self as u64
            }
            fn from_u64(n: u64) -> Self {
                n as _
            }
            #[cfg(debug_assertions)]
            fn as_u128(&self) -> u128 {
                *self as u128
            }
            #[cfg(debug_assertions)]
            fn from_u128(n: u64) -> Self {
                n as _
            }
        }
    };
}
macro_rules! double_int {
    ( $x:ty, $y:ty ) => {
        int!($x);
        impl DoubleInt for $x {
            type Double = $y;

            fn as_double(self) -> $y {
                self as _
            }
            fn from_double(n: $y) -> $x {
                n as _
            }
        }
    };
}

double_int!(u32, u64);
double_int!(u64, u128);
int!(u128);

// extended Euclid algorithm
// precondition: a is odd
pub(crate) fn modular_inverse<T: Int>(a: T) -> T {
    let zero = T::zero();
    let one = T::one();
    debug_assert!(a % (one + one) == one, "{:?} is not odd", a);

    let mut t = zero;
    let mut newt = one;
    let mut r = zero;
    let mut newr = a;

    while newr != zero {
        let quot = if r == zero {
            // special case when we're just starting out
            // This works because we know that
            // a does not divide 2^64, so floor(2^64 / a) == floor((2^64-1) / a);
            T::max_value()
        } else {
            r
        } / newr;

        let newtp = t.wrapping_sub(&quot.wrapping_mul(&newt));
        t = newt;
        newt = newtp;

        let newrp = r.wrapping_sub(&quot.wrapping_mul(&newr));
        r = newr;
        newr = newrp;
    }

    debug_assert_eq!(r, one);
    t
}

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

    macro_rules! parametrized_check {
        ( $f:ident ) => {
            paste::item! {
                #[test]
                fn [< $f _ u32 >]() {
                    $f::<u32>()
                }
                #[test]
                fn [< $f _ u64 >]() {
                    $f::<u64>()
                }
            }
        };
    }

    fn test_inverter<T: Int>() {
        // All odd integers from 1 to 20 000
        let one = T::from(1).unwrap();
        let two = T::from(2).unwrap();
        let mut test_values = (0..10_000)
            .map(|i| T::from(i).unwrap())
            .map(|i| two * i + one);

        assert!(test_values.all(|x| x.wrapping_mul(&modular_inverse(x)) == one));
    }
    parametrized_check!(test_inverter);

    fn test_add<A: DoubleInt>() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = Montgomery::<A>::new(n);
            for x in 0..n {
                let m_x = m.from_u64(x);
                for y in 0..=x {
                    let m_y = m.from_u64(y);
                    println!("{n:?}, {x:?}, {y:?}", n = n, x = x, y = y);
                    assert_eq!((x + y) % n, m.to_u64(m.add(m_x, m_y)));
                }
            }
        }
    }
    parametrized_check!(test_add);

    fn test_mult<A: DoubleInt>() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = Montgomery::<A>::new(n);
            for x in 0..n {
                let m_x = m.from_u64(x);
                for y in 0..=x {
                    let m_y = m.from_u64(y);
                    assert_eq!((x * y) % n, m.to_u64(m.mul(m_x, m_y)));
                }
            }
        }
    }
    parametrized_check!(test_mult);

    fn test_roundtrip<A: DoubleInt>() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = Montgomery::<A>::new(n);
            for x in 0..n {
                let x_ = m.from_u64(x);
                assert_eq!(x, m.to_u64(x_));
            }
        }
    }
    parametrized_check!(test_roundtrip);
}
<commit_msg>factor::numeric::DoubleInt: Clarify methods and associated types<commit_after>// * This file is part of the uutils coreutils package.
// *
// * (c) 2015 Wiktor Kuropatwa <wiktor.kuropatwa@gmail.com>
// * (c) 2020 nicoo            <nicoo@debian.org>
// *
// * For the full copyright and license information, please view the LICENSE file
// * that was distributed with this source code.

use num_traits::{
    identities::{One, Zero},
    int::PrimInt,
    ops::wrapping::{WrappingMul, WrappingNeg, WrappingSub},
};
use std::fmt::{Debug, Display};
use std::mem::swap;

// This is incorrectly reported as dead code,
//  presumably when included in build.rs.
#[allow(dead_code)]
pub(crate) fn gcd(mut a: u64, mut b: u64) -> u64 {
    while b > 0 {
        a %= b;
        swap(&mut a, &mut b);
    }
    a
}

pub(crate) trait Arithmetic: Copy + Sized {
    // The type of integers mod m, in some opaque representation
    type ModInt: Copy + Sized + Eq;

    fn new(m: u64) -> Self;
    fn modulus(&self) -> u64;
    fn from_u64(&self, n: u64) -> Self::ModInt;
    fn to_u64(&self, n: Self::ModInt) -> u64;
    fn add(&self, a: Self::ModInt, b: Self::ModInt) -> Self::ModInt;
    fn mul(&self, a: Self::ModInt, b: Self::ModInt) -> Self::ModInt;

    fn pow(&self, mut a: Self::ModInt, mut b: u64) -> Self::ModInt {
        let (_a, _b) = (a, b);
        let mut result = self.one();
        while b > 0 {
            if b & 1 != 0 {
                result = self.mul(result, a);
            }
            a = self.mul(a, a);
            b >>= 1;
        }

        // Check that r (reduced back to the usual representation) equals
        //  a^b % n, unless the latter computation overflows
        // Temporarily commented-out, as there u64::checked_pow is not available
        //  on the minimum supported Rust version, nor is an appropriate method
        //  for compiling the check conditionally.
        //debug_assert!(self
        //    .to_u64(_a)
        //    .checked_pow(_b as u32)
        //    .map(|r| r % self.modulus() == self.to_u64(result))
        //    .unwrap_or(true));

        result
    }

    fn one(&self) -> Self::ModInt {
        self.from_u64(1)
    }
    fn minus_one(&self) -> Self::ModInt {
        self.from_u64(self.modulus() - 1)
    }
    fn zero(&self) -> Self::ModInt {
        self.from_u64(0)
    }
}

#[derive(Clone, Copy, Debug)]
pub(crate) struct Montgomery<T: DoubleInt> {
    a: T,
    n: T,
}

impl<T: DoubleInt> Montgomery<T> {
    /// computes x/R mod n efficiently
    fn reduce(&self, x: T::DoubleWidth) -> T {
        let t_bits = T::zero().count_zeros() as usize;

        debug_assert!(x < (self.n.as_double_width()) << t_bits);
        // TODO: optimiiiiiiise
        let Montgomery { a, n } = self;
        let m = T::from_double_width(x).wrapping_mul(a);
        let nm = (n.as_double_width()) * (m.as_double_width());
        let (xnm, overflow) = x.overflowing_add_(nm); // x + n*m
        debug_assert_eq!(
            xnm % (T::DoubleWidth::one() << T::zero().count_zeros() as usize),
            T::DoubleWidth::zero()
        );

        // (x + n*m) / R
        // in case of overflow, this is (2¹²⁸ + xnm)/2⁶⁴ - n = xnm/2⁶⁴ + (2⁶⁴ - n)
        let y = T::from_double_width(xnm >> t_bits)
            + if !overflow {
                T::zero()
            } else {
                n.wrapping_neg()
            };

        if y >= *n {
            y - *n
        } else {
            y
        }
    }
}

impl<T: DoubleInt> Arithmetic for Montgomery<T> {
    // Montgomery transform, R=2⁶⁴
    // Provides fast arithmetic mod n (n odd, u64)
    type ModInt = T;

    fn new(n: u64) -> Self {
        debug_assert!(T::zero().count_zeros() >= 64 || n < (1 << T::zero().count_zeros() as usize));
        let n = T::from_u64(n);
        let a = modular_inverse(n).wrapping_neg();
        debug_assert_eq!(n.wrapping_mul(&a), T::one().wrapping_neg());
        Montgomery { a, n }
    }

    fn modulus(&self) -> u64 {
        self.n.as_u64()
    }

    fn from_u64(&self, x: u64) -> Self::ModInt {
        // TODO: optimise!
        debug_assert!(x < self.n.as_u64());
        let r = T::from_double_width(
            ((T::DoubleWidth::from_u64(x)) << T::zero().count_zeros() as usize)
                % self.n.as_double_width(),
        );
        debug_assert_eq!(x, self.to_u64(r));
        r
    }

    fn to_u64(&self, n: Self::ModInt) -> u64 {
        self.reduce(n.as_double_width()).as_u64()
    }

    fn add(&self, a: Self::ModInt, b: Self::ModInt) -> Self::ModInt {
        let (r, overflow) = a.overflowing_add_(b);

        // In case of overflow, a+b = 2⁶⁴ + r = (2⁶⁴ - n) + r (working mod n)
        let r = if !overflow {
            r
        } else {
            r + self.n.wrapping_neg()
        };

        // Normalise to [0; n[
        let r = if r < self.n { r } else { r - self.n };

        // Check that r (reduced back to the usual representation) equals
        // a+b % n
        #[cfg(debug_assertions)]
        {
            let a_r = self.to_u64(a) as u128;
            let b_r = self.to_u64(b) as u128;
            let r_r = self.to_u64(r);
            let r_2 = ((a_r + b_r) % self.n.as_u128()) as u64;
            debug_assert_eq!(
                r_r, r_2,
                "[{}] = {} ≠ {} = {} + {} = [{}] + [{}] mod {}; a = {}",
                r, r_r, r_2, a_r, b_r, a, b, self.n, self.a
            );
        }
        r
    }

    fn mul(&self, a: Self::ModInt, b: Self::ModInt) -> Self::ModInt {
        let r = self.reduce(a.as_double_width() * b.as_double_width());

        // Check that r (reduced back to the usual representation) equals
        // a*b % n
        #[cfg(debug_assertions)]
        {
            let a_r = self.to_u64(a) as u128;
            let b_r = self.to_u64(b) as u128;
            let r_r = self.to_u64(r);
            let r_2: u64 = ((a_r * b_r) % self.n.as_u128()) as u64;
            debug_assert_eq!(
                r_r, r_2,
                "[{}] = {} ≠ {} = {} * {} = [{}] * [{}] mod {}; a = {}",
                r, r_r, r_2, a_r, b_r, a, b, self.n, self.a
            );
        }
        r
    }
}

// NOTE: Trait can be removed once num-traits adds a similar one;
//       see https://github.com/rust-num/num-traits/issues/168
pub(crate) trait OverflowingAdd: Sized {
    fn overflowing_add_(self, n: Self) -> (Self, bool);
}
macro_rules! overflowing {
    ($x:ty) => {
        impl OverflowingAdd for $x {
            fn overflowing_add_(self, n: Self) -> (Self, bool) {
                self.overflowing_add(n)
            }
        }
    };
}
overflowing!(u32);
overflowing!(u64);
overflowing!(u128);

pub(crate) trait Int:
    Display + Debug + PrimInt + OverflowingAdd + WrappingNeg + WrappingSub + WrappingMul
{
    fn as_u64(&self) -> u64;
    fn from_u64(n: u64) -> Self;
    #[cfg(debug_assertions)]
    fn as_u128(&self) -> u128;
    #[cfg(debug_assertions)]
    fn from_u128(n: u64) -> Self;
}

pub(crate) trait DoubleInt: Int {
    type DoubleWidth: Int;

    fn as_double_width(self) -> Self::DoubleWidth;
    fn from_double_width(n: Self::DoubleWidth) -> Self;
}

macro_rules! int {
    ( $x:ty ) => {
        impl Int for $x {
            fn as_u64(&self) -> u64 {
                *self as u64
            }
            fn from_u64(n: u64) -> Self {
                n as _
            }
            #[cfg(debug_assertions)]
            fn as_u128(&self) -> u128 {
                *self as u128
            }
            #[cfg(debug_assertions)]
            fn from_u128(n: u64) -> Self {
                n as _
            }
        }
    };
}
macro_rules! double_int {
    ( $x:ty, $y:ty ) => {
        int!($x);
        impl DoubleInt for $x {
            type DoubleWidth = $y;

            fn as_double_width(self) -> $y {
                self as _
            }
            fn from_double_width(n: $y) -> $x {
                n as _
            }
        }
    };
}

double_int!(u32, u64);
double_int!(u64, u128);
int!(u128);

// extended Euclid algorithm
// precondition: a is odd
pub(crate) fn modular_inverse<T: Int>(a: T) -> T {
    let zero = T::zero();
    let one = T::one();
    debug_assert!(a % (one + one) == one, "{:?} is not odd", a);

    let mut t = zero;
    let mut newt = one;
    let mut r = zero;
    let mut newr = a;

    while newr != zero {
        let quot = if r == zero {
            // special case when we're just starting out
            // This works because we know that
            // a does not divide 2^64, so floor(2^64 / a) == floor((2^64-1) / a);
            T::max_value()
        } else {
            r
        } / newr;

        let newtp = t.wrapping_sub(&quot.wrapping_mul(&newt));
        t = newt;
        newt = newtp;

        let newrp = r.wrapping_sub(&quot.wrapping_mul(&newr));
        r = newr;
        newr = newrp;
    }

    debug_assert_eq!(r, one);
    t
}

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

    macro_rules! parametrized_check {
        ( $f:ident ) => {
            paste::item! {
                #[test]
                fn [< $f _ u32 >]() {
                    $f::<u32>()
                }
                #[test]
                fn [< $f _ u64 >]() {
                    $f::<u64>()
                }
            }
        };
    }

    fn test_inverter<T: Int>() {
        // All odd integers from 1 to 20 000
        let one = T::from(1).unwrap();
        let two = T::from(2).unwrap();
        let mut test_values = (0..10_000)
            .map(|i| T::from(i).unwrap())
            .map(|i| two * i + one);

        assert!(test_values.all(|x| x.wrapping_mul(&modular_inverse(x)) == one));
    }
    parametrized_check!(test_inverter);

    fn test_add<A: DoubleInt>() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = Montgomery::<A>::new(n);
            for x in 0..n {
                let m_x = m.from_u64(x);
                for y in 0..=x {
                    let m_y = m.from_u64(y);
                    println!("{n:?}, {x:?}, {y:?}", n = n, x = x, y = y);
                    assert_eq!((x + y) % n, m.to_u64(m.add(m_x, m_y)));
                }
            }
        }
    }
    parametrized_check!(test_add);

    fn test_mult<A: DoubleInt>() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = Montgomery::<A>::new(n);
            for x in 0..n {
                let m_x = m.from_u64(x);
                for y in 0..=x {
                    let m_y = m.from_u64(y);
                    assert_eq!((x * y) % n, m.to_u64(m.mul(m_x, m_y)));
                }
            }
        }
    }
    parametrized_check!(test_mult);

    fn test_roundtrip<A: DoubleInt>() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = Montgomery::<A>::new(n);
            for x in 0..n {
                let x_ = m.from_u64(x);
                assert_eq!(x, m.to_u64(x_));
            }
        }
    }
    parametrized_check!(test_roundtrip);
}
<|endoftext|>


<commit_before>use crate::display_delegate::display;
use proc_macro2::{Group, Span, TokenTree};
use quote::ToTokens;
use std::fmt::{Error, Formatter};
use syn::punctuated::Punctuated;
use syn::token::{Colon2, Semi};
use syn::{
    self, ArgCaptured, Block, Expr, ExprVerbatim, FnArg, FnDecl, GenericParam, Ident, Pat,
    PatIdent, PathSegment, Stmt,
};

const MOCKTOPUS_CRATE_NAME: &str = "__mocktopus_crate__";
const STD_CRATE_NAME: &str = "__mocktopus_std__";
const ARGS_TO_CONTINUE_NAME: &str = "__mocktopus_args_to_continue__";
const UNWIND_DATA_NAME: &str = "__mocktopus_unwind_data__";

macro_rules! error_msg {
    ($msg:expr) => {
        concat!("Mocktopus internal error: ", $msg)
    };
}

pub enum FnHeaderBuilder<'a> {
    StaticFn,
    StructImpl,
    TraitDefault,
    TraitImpl(&'a Punctuated<PathSegment, Colon2>),
}

impl<'a> FnHeaderBuilder<'a> {
    pub fn build(&self, fn_ident: &Ident, fn_decl: &FnDecl, fn_block_span: Span) -> Stmt {
        let fn_args = &fn_decl.inputs;
        let header_str = format!(
            r#"{{
    extern crate mocktopus as {mocktopus};
    extern crate std as {std_crate};
    match {std_crate}::panic::catch_unwind({std_crate}::panic::AssertUnwindSafe (
            || {mocktopus}::mocking::Mockable::call_mock(&{full_fn_name}, {extract_args}))) {{
        Ok({mocktopus}::mocking::MockResult::Continue({args_to_continue})) => {restore_args},
        Ok({mocktopus}::mocking::MockResult::Return(result)) => {{
            {forget_args}
            let returned = unsafe {{ {std_crate}::mem::transmute_copy(&result) }};
            {std_crate}::mem::forget(result);
            return returned;
        }},
        Err({unwind}) => {{
            {forget_args}
            {std_crate}::panic::resume_unwind({unwind});
        }},
    }}
}}"#,
            mocktopus = MOCKTOPUS_CRATE_NAME,
            std_crate = STD_CRATE_NAME,
            full_fn_name = display(|f| write_full_fn_name(f, self, fn_ident, fn_decl)),
            extract_args = display(|f| write_extract_args(f, fn_args)),
            args_to_continue = ARGS_TO_CONTINUE_NAME,
            restore_args = display(|f| write_restore_args(f, fn_args)),
            forget_args = display(|f| write_forget_args(f, fn_args)),
            unwind = UNWIND_DATA_NAME
        );
        let header_block =
            syn::parse_str::<Block>(&header_str).expect(error_msg!("generated header unparsable"));
        create_call_site_spanned_stmt(header_block, fn_block_span)
    }
}

fn create_call_site_spanned_stmt(block: Block, span: Span) -> Stmt {
    let token_stream = block
        .into_token_stream()
        .into_iter()
        .map(|tt| make_token_tree_span_call_site(tt, span))
        .collect();
    Stmt::Semi(
        Expr::Verbatim(ExprVerbatim { tts: token_stream }),
        Semi([span]),
    )
}

fn make_token_tree_span_call_site(mut token_tree: TokenTree, span: Span) -> TokenTree {
    token_tree.set_span(span);
    if let TokenTree::Group(ref mut group) = token_tree {
        let tokens = group
            .stream()
            .into_iter()
            .map(|tt| make_token_tree_span_call_site(tt, span))
            .collect();
        *group = Group::new(group.delimiter(), tokens);
    }
    token_tree
}

fn write_full_fn_name(
    f: &mut Formatter,
    builder: &FnHeaderBuilder,
    fn_ident: &Ident,
    fn_decl: &FnDecl,
) -> Result<(), Error> {
    match *builder {
        FnHeaderBuilder::StaticFn => (),
        FnHeaderBuilder::StructImpl | FnHeaderBuilder::TraitDefault => write!(f, "Self::")?,
        FnHeaderBuilder::TraitImpl(ref path) => {
            write!(f, "<Self as {}>::", display(|f| write_trait_path(f, path)))?
        }
    }
    write!(
        f,
        "{}::<{}>",
        fn_ident,
        display(|f| write_fn_generics(f, fn_decl))
    )
}

fn write_trait_path<T: ToTokens + Clone>(
    f: &mut Formatter,
    path: &Punctuated<PathSegment, T>,
) -> Result<(), Error> {
    write!(f, "{}", path.into_token_stream())
}

fn write_fn_generics(f: &mut Formatter, fn_decl: &FnDecl) -> Result<(), Error> {
    fn_decl
        .generics
        .params
        .iter()
        .filter_map(get_generic_param_name)
        .map(|param| write!(f, "{},", param))
        .collect()
}

fn get_generic_param_name(param: &GenericParam) -> Option<String> {
    match *param {
        GenericParam::Type(ref type_param) => Some(type_param.ident.to_string()),
        _ => None,
    }
}

fn write_extract_args<T>(f: &mut Formatter, fn_args: &Punctuated<FnArg, T>) -> Result<(), Error> {
    if fn_args.is_empty() {
        return write!(f, "()");
    }
    write!(f, "unsafe {{ (")?;
    for fn_arg_name in iter_fn_arg_names(fn_args) {
        write!(
            f,
            "{}::mem::transmute_copy(&{}), ",
            STD_CRATE_NAME, fn_arg_name
        )?;
    }
    write!(f, ") }}")
}

fn write_restore_args<T>(f: &mut Formatter, fn_args: &Punctuated<FnArg, T>) -> Result<(), Error> {
    if fn_args.is_empty() {
        return writeln!(f, "()");
    }
    writeln!(f, "unsafe {{")?;
    for (fn_arg_index, fn_arg_name) in iter_fn_arg_names(fn_args).enumerate() {
        writeln!(
            f,
            "{}::mem::forget({}::mem::replace({}::mocking_utils::as_mut(&{}), {}.{}));",
            STD_CRATE_NAME,
            STD_CRATE_NAME,
            MOCKTOPUS_CRATE_NAME,
            fn_arg_name,
            ARGS_TO_CONTINUE_NAME,
            fn_arg_index
        )?;
    }
    writeln!(f, "}}")
}

fn write_forget_args<T>(f: &mut Formatter, fn_args: &Punctuated<FnArg, T>) -> Result<(), Error> {
    for fn_arg_name in iter_fn_arg_names(fn_args) {
        writeln!(f, "{}::mem::forget({});", STD_CRATE_NAME, fn_arg_name)?;
    }
    Ok(())
}

fn iter_fn_arg_names<'a, T>(
    input_args: &'a Punctuated<FnArg, T>,
) -> impl Iterator<Item = String> + 'a {
    input_args.iter().map(|fn_arg| match *fn_arg {
        FnArg::SelfRef(_) | FnArg::SelfValue(_) => "self".to_string(),
        FnArg::Captured(ArgCaptured {
            pat:
                Pat::Ident(PatIdent {
                    ref ident,
                    subpat: None,
                    ..
                }),
            ..
        }) => ident.to_string(),
        _ => panic!(
            "{}: '{}'",
            error_msg!("invalid fn arg type"),
            fn_arg.clone().into_token_stream()
        ),
    })
}
<commit_msg>Clean up and name collision proof macro generated code<commit_after>use crate::display_delegate::display;
use proc_macro2::{Group, Span, TokenTree};
use quote::ToTokens;
use std::fmt::{Error, Formatter};
use syn::punctuated::Punctuated;
use syn::token::{Colon2, Semi};
use syn::{
    self, ArgCaptured, Block, Expr, ExprVerbatim, FnArg, FnDecl, GenericParam, Ident, Pat,
    PatIdent, PathSegment, Stmt,
};

const MOCKTOPUS_CRATE_NAME: &str = "__mocktopus_crate__";
const STD_CRATE_NAME: &str = "__mocktopus_std__";
const ARGS_TO_CONTINUE_NAME: &str = "__mocktopus_args_to_continue__";
const ARGS_TO_RETURN_NAME: &str = "__mocktopus_args_to_return__";
const UNWIND_DATA_NAME: &str = "__mocktopus_unwind_data__";

macro_rules! error_msg {
    ($msg:expr) => {
        concat!("Mocktopus internal error: ", $msg)
    };
}

pub enum FnHeaderBuilder<'a> {
    StaticFn,
    StructImpl,
    TraitDefault,
    TraitImpl(&'a Punctuated<PathSegment, Colon2>),
}

impl<'a> FnHeaderBuilder<'a> {
    pub fn build(&self, fn_ident: &Ident, fn_decl: &FnDecl, fn_block_span: Span) -> Stmt {
        let fn_args = &fn_decl.inputs;
        let header_str = format!(
            r#"{{
    extern crate mocktopus as {mocktopus};
    extern crate std as {std_crate};
    match {std_crate}::panic::catch_unwind({std_crate}::panic::AssertUnwindSafe (
            || {mocktopus}::mocking::Mockable::call_mock(&{full_fn_name}, {extract_args}))) {{
        Ok({mocktopus}::mocking::MockResult::Continue(mut {args_to_continue})) => {restore_args},
        Ok({mocktopus}::mocking::MockResult::Return({args_to_return})) => {{
            {forget_args}
            let returned = unsafe {{ {std_crate}::mem::transmute_copy(&{args_to_return}) }};
            {std_crate}::mem::forget({args_to_return});
            return returned;
        }},
        Err({unwind}) => {{
            {forget_args}
            {std_crate}::panic::resume_unwind({unwind});
        }},
    }}
}}"#,
            mocktopus = MOCKTOPUS_CRATE_NAME,
            std_crate = STD_CRATE_NAME,
            full_fn_name = display(|f| write_full_fn_name(f, self, fn_ident, fn_decl)),
            extract_args = display(|f| write_extract_args(f, fn_args)),
            args_to_continue = ARGS_TO_CONTINUE_NAME,
            args_to_return = ARGS_TO_RETURN_NAME,
            restore_args = display(|f| write_restore_args(f, fn_args)),
            forget_args = display(|f| write_forget_args(f, fn_args)),
            unwind = UNWIND_DATA_NAME
        );
        let header_block =
            syn::parse_str::<Block>(&header_str).expect(error_msg!("generated header unparsable"));
        create_call_site_spanned_stmt(header_block, fn_block_span)
    }
}

fn create_call_site_spanned_stmt(block: Block, span: Span) -> Stmt {
    let token_stream = block
        .into_token_stream()
        .into_iter()
        .map(|tt| make_token_tree_span_call_site(tt, span))
        .collect();
    Stmt::Semi(
        Expr::Verbatim(ExprVerbatim { tts: token_stream }),
        Semi([span]),
    )
}

fn make_token_tree_span_call_site(mut token_tree: TokenTree, span: Span) -> TokenTree {
    token_tree.set_span(span);
    if let TokenTree::Group(ref mut group) = token_tree {
        let tokens = group
            .stream()
            .into_iter()
            .map(|tt| make_token_tree_span_call_site(tt, span))
            .collect();
        *group = Group::new(group.delimiter(), tokens);
    }
    token_tree
}

fn write_full_fn_name(
    f: &mut Formatter,
    builder: &FnHeaderBuilder,
    fn_ident: &Ident,
    fn_decl: &FnDecl,
) -> Result<(), Error> {
    match *builder {
        FnHeaderBuilder::StaticFn => (),
        FnHeaderBuilder::StructImpl | FnHeaderBuilder::TraitDefault => write!(f, "Self::")?,
        FnHeaderBuilder::TraitImpl(ref path) => {
            write!(f, "<Self as {}>::", display(|f| write_trait_path(f, path)))?
        }
    }
    write!(
        f,
        "{}::<{}>",
        fn_ident,
        display(|f| write_fn_generics(f, fn_decl))
    )
}

fn write_trait_path<T: ToTokens + Clone>(
    f: &mut Formatter,
    path: &Punctuated<PathSegment, T>,
) -> Result<(), Error> {
    write!(f, "{}", path.into_token_stream())
}

fn write_fn_generics(f: &mut Formatter, fn_decl: &FnDecl) -> Result<(), Error> {
    fn_decl
        .generics
        .params
        .iter()
        .filter_map(get_generic_param_name)
        .map(|param| write!(f, "{},", param))
        .collect()
}

fn get_generic_param_name(param: &GenericParam) -> Option<String> {
    match *param {
        GenericParam::Type(ref type_param) => Some(type_param.ident.to_string()),
        _ => None,
    }
}

fn write_extract_args<T>(f: &mut Formatter, fn_args: &Punctuated<FnArg, T>) -> Result<(), Error> {
    if fn_args.is_empty() {
        return write!(f, "()");
    }
    write!(f, "unsafe {{ (")?;
    for fn_arg_name in iter_fn_arg_names(fn_args) {
        write!(
            f,
            "{}::mem::transmute_copy(&{}), ",
            STD_CRATE_NAME, fn_arg_name
        )?;
    }
    write!(f, ") }}")
}

fn write_restore_args<T>(f: &mut Formatter, fn_args: &Punctuated<FnArg, T>) -> Result<(), Error> {
    if fn_args.is_empty() {
        return writeln!(f, "()");
    }
    writeln!(f, "unsafe {{")?;
    for (fn_arg_index, fn_arg_name) in iter_fn_arg_names(fn_args).enumerate() {
        writeln!(
            f,
            "{}::mem::swap(&mut *(&{} as *const _ as *mut _), &mut {}.{});",
            STD_CRATE_NAME,
            fn_arg_name,
            ARGS_TO_CONTINUE_NAME,
            fn_arg_index
        )?;
    }
    writeln!(f, "{}::mem::forget({});", STD_CRATE_NAME, ARGS_TO_CONTINUE_NAME)?;
    writeln!(f, "}}")
}

fn write_forget_args<T>(f: &mut Formatter, fn_args: &Punctuated<FnArg, T>) -> Result<(), Error> {
    for fn_arg_name in iter_fn_arg_names(fn_args) {
        writeln!(f, "{}::mem::forget({});", STD_CRATE_NAME, fn_arg_name)?;
    }
    Ok(())
}

fn iter_fn_arg_names<'a, T>(
    input_args: &'a Punctuated<FnArg, T>,
) -> impl Iterator<Item = String> + 'a {
    input_args.iter().map(|fn_arg| match *fn_arg {
        FnArg::SelfRef(_) | FnArg::SelfValue(_) => "self".to_string(),
        FnArg::Captured(ArgCaptured {
            pat:
                Pat::Ident(PatIdent {
                    ref ident,
                    subpat: None,
                    ..
                }),
            ..
        }) => ident.to_string(),
        _ => panic!(
            "{}: '{}'",
            error_msg!("invalid fn arg type"),
            fn_arg.clone().into_token_stream()
        ),
    })
}
<|endoftext|>


<commit_before>//! Generate the site content.

// Standard library
use std::borrow::Cow::{Borrowed, Owned};
use std::fs::File;
use std::io::prelude::*;
use std::path::{Path, PathBuf};

// Third party
use glob::{glob, Paths};
use pulldown_cmark as cmark;
use pulldown_cmark::{Event, Parser, Tag};
use syntect::highlighting::ThemeSet;

// First party
use crate::config::{Config, SyntaxOption};

/// Load the `Paths` for all markdown files in the specified content directory.
///
/// If the specified configuration values are wrong, we return a vector of all
/// the errors in the config as far as this is concerned.
///
/// TODO: move that validation elsewhere? Should probably extract a more general
/// site configuration validator.
fn load_markdown_paths(site_directory: &PathBuf, config: &Config) -> Result<Vec<PathBuf>, String> {
    let site = site_directory
        .to_str()
        .ok_or(format!("bad `site`: {:?}", site_directory));

    let directories = config.directories.content.to_str().ok_or(format!(
        "bad content directory: {:?}",
        config.directories.content
    ));

    let (site, directories) = match (site, directories) {
        (Err(s), Err(d)) => Err(s + "\n" + &d),
        (Err(s), Ok(_)) => Err(s),
        (Ok(_), Err(d)) => Err(d),
        (Ok(s), Ok(d)) => Ok((s, d)),
    }?;

    let content_glob = format!("{}/{}/**/*.md", site, directories);

    let paths = glob(&content_glob).map_err(|err| format!("{:?}", err))?;

    let (valid, errors): (Vec<Result<PathBuf, _>>, Vec<Result<_, String>>) = paths
        .map(|glob_result| glob_result.map_err(|err| format!("{:?}", err)))
        .partition(|path| path.is_ok());

    let errors: Vec<String> = errors.into_iter().map(|e| e.unwrap_err()).collect();
    if errors.len() == 0 {
        Ok(valid.into_iter().map(|v| v.unwrap()).collect())
    } else {
        Err(errors.join(",\n"))
    }
}

// TODO: this is an obvious candidate for a `HashMap`, so use that instead.
type LoadTuple<'a> = (&'a PathBuf, Result<String, String>);

// Instead of using `.collect` at the end of these, I *should* be able to use
// something like `Result<impl Iterator<HashMap<PathBuf, String>, String>`, even
// if the type isn't *quite* taht.
fn load_content<'p>(paths: &'p [PathBuf]) -> Result<Vec<(&'p PathBuf, String)>, String> {
    let (contents, errs): (Vec<LoadTuple>, Vec<LoadTuple>) = paths
        .into_iter()
        .map(|path| {
            let content = std::fs::read_to_string(path).map_err(|e| format!("{:?}", e));
            (path, content)
        })
        .partition(|(_, result)| result.is_ok());

    if errs.len() == 0 {
        Ok(contents
            .into_iter()
            .map(|(path, contents)| (path, contents.unwrap()))
            .collect())
    } else {
        Err(errs
            .into_iter()
            .map(|(_path, err)| err.unwrap_err())
            .collect::<Vec<String>>()
            .join(",\n"))
    }
}

/// Load the templates associated with each taxonomy.
fn load_templates(_site_directory: &PathBuf, _config: &Config) -> Result<Paths, String> {
    unimplemented!()
}

// TODO: put this somewhere else when extracting all the parser logic below into
// its own module.
#[derive(Debug)]
enum ParseState<'s> {
    CodeBlock(&'s syntect::parsing::SyntaxReference),
    PlainTextBlock,
    NonCode,
}

impl<'s> std::default::Default for ParseState<'s> {
    fn default() -> ParseState<'s> {
        ParseState::NonCode
    }
}

enum Syntax {
    Highlight(syntect::highlighting::Theme),
    TagOnly,
    Off,
}

impl Syntax {
    fn try_from(config_syntax: &SyntaxOption) -> Result<Syntax, String> {
        match config_syntax {
            SyntaxOption::Highlight(theme) => {
                let theme_file = PathBuf::from("data").join(theme);
                let theme = ThemeSet::get_theme(theme_file).map_err(|err| format!("{:?}", err))?;
                Ok(Syntax::Highlight(theme))
            }
            SyntaxOption::TagOnly => Ok(Syntax::TagOnly),
            SyntaxOption::Off => Ok(Syntax::Off),
        }
    }
}

/// Generate content from a configuration.
pub fn build(site_directory: PathBuf) -> Result<(), String> {
    // NOTE: this is almost certainly not what we *ultimately* want here, but
    // it'll do as a starting point. There's a lot of `into_iter` and `collect`
    // happening here, which forces us to eagerly materialize a lot of these
    // collections rather than being able to deal with them in a more
    // "streaming" fashion.
    let config = Config::from_file(&PathBuf::from(&site_directory))?;
    let markdown_paths = load_markdown_paths(&site_directory, &config)?;
    let contents = load_content(&markdown_paths)?;

    let configured_syntax = Syntax::try_from(&config.options.syntax)?;

    // TODO: generate these! Best move: on build, generate a `.packdump`, since
    // that seems to be what Syntect knows how to deal with.
    let syntax_dir = PathBuf::from("data/syntaxes");
    let syntax_set =
        syntect::parsing::SyntaxSet::load_from_folder(syntax_dir).expect("can't load syntaxes");

    let parsed_content: Vec<(&PathBuf, String)> = contents
        .into_iter()
        .map(|(path, content)| {
            let parser = Parser::new_ext(&content, pulldown_cmark::Options::all()).scan(
                ParseState::default(),
                |state, event| match event {
                    Event::Start(Tag::CodeBlock(ref language)) => match state {
                        ParseState::NonCode => {
                            *state = match syntax_set.find_syntax_by_name(&language) {
                                Some(syntax) => ParseState::CodeBlock(syntax),
                                None => ParseState::PlainTextBlock,
                            };

                            // We start every code block with `<pre><code class="...">`
                            // so that we always have semantically correct HTML. However, we
                            // don't add the language class if the language isn't set.
                            let content = if language.len() > 0 {
                                Owned(format!(r#"
"#, language))
                            } else {
                                Borrowed("
")
                            };

                            Some(Event::Html(content))
                        }
                        ParseState::CodeBlock(_) | ParseState::PlainTextBlock => {
                            unreachable!("Bad event/state: {:?} with {:?}", event, state);
                        }
                    },

                    Event::End(Tag::CodeBlock(_)) => match state {
                        ParseState::CodeBlock(_) | ParseState::PlainTextBlock => {
                            *state = ParseState::NonCode;
                            Some(Event::Html(Borrowed("
")))
                        }
                        ParseState::NonCode => {
                            unreachable!("Bad event/state: {:?} with {:?}", event, state);
                        }
                    },

                    // When we are in a text block, we *may* be in a code block,
                    // so we may also need to do whatever highlighting is
                    // specified.
                    Event::Text(s) => {
                        match state {
                            ParseState::CodeBlock(syntax) => match &configured_syntax {
                                Syntax::Highlight(theme) => {
                                    let highlighted = syntect::html::highlighted_html_for_string(
                                        &s,
                                        &syntax_set,
                                        &syntax,
                                        theme,
                                    );

                                    Some(Event::Text(Owned(highlighted)))
                                }
                                Syntax::TagOnly => {
                                    // TODO: implement tag-only parsing with syntect
                                    unimplemented!()
                                }
                                Syntax::Off => Some(Event::Text(s)),
                            },
                            ParseState::PlainTextBlock => Some(Event::Text(s)),
                            ParseState::NonCode => Some(Event::Text(s)),
                        }
                    }

                    // TODO: accumulate footnote references and always put them
                    // at the end. Also, generate back-links for them, unless or
                    // until pulldown-cmark does so.
                    Event::FootnoteReference(s) => Some(Event::FootnoteReference(s)),

                    Event::Start(s) => Some(Event::Start(s)),
                    Event::End(e) => Some(Event::End(e)),
                    Event::InlineHtml(s) => Some(Event::InlineHtml(s)),
                    Event::Html(h) => Some(Event::Html(h)),
                    Event::SoftBreak => Some(Event::SoftBreak),
                    Event::HardBreak => Some(Event::HardBreak),
                },
            );

            // These numbers are derived from a brief survey of content in my
            // own website historically, including both code-heavy and
            // prose-heavy posts. It's probably worth exploring other corpora
            // to see if this is basically right.
            let length_bound_estimate = match &config.options.syntax {
                SyntaxOption::Off | SyntaxOption::TagOnly => 2,
                SyntaxOption::Highlight(_) => 8, // TODO: check other pages
            };

            let mut html = String::with_capacity(content.len() * length_bound_estimate);
            cmark::html::push_html(&mut html, parser);

            (path, html)
        })
        .collect();

    // TODO: extract this as part of the writing it out process.
    let output_dir = site_directory.join(&config.directories.output);

    if output_dir.exists() && !output_dir.is_dir() {
        return Err(format!("Output directory {:?} already exists!", output_dir));
    }

    if !output_dir.exists() {
        std::fs::create_dir(&output_dir).expect(&format!(
            "could not create output directory {:?}",
            &output_dir
        ));
    }

    for (path, content) in parsed_content {
        let file_path = path
            .file_name()
            .ok_or(format!("invalid file: {:?}", path))?;

        let dest = output_dir.join(file_path).with_extension("html");

        let mut fd = match File::create(&dest) {
            Ok(file) => file,
            Err(reason) => {
                return Err(format!(
                    "Could not open {} for write: {}",
                    dest.to_string_lossy(),
                    reason
                ));
            }
        };

        let result = write!(fd, "{}", content);
        if let Err(reason) = result {
            return Err(format!("{:?}", reason.kind()));
        }
    }

    Ok(())
}
<commit_msg>fix: better error message for output-dir-exists.<commit_after>//! Generate the site content.

// Standard library
use std::borrow::Cow::{Borrowed, Owned};
use std::fs::File;
use std::io::prelude::*;
use std::path::{Path, PathBuf};

// Third party
use glob::{glob, Paths};
use pulldown_cmark as cmark;
use pulldown_cmark::{Event, Parser, Tag};
use syntect::highlighting::ThemeSet;

// First party
use crate::config::{Config, SyntaxOption};

/// Load the `Paths` for all markdown files in the specified content directory.
///
/// If the specified configuration values are wrong, we return a vector of all
/// the errors in the config as far as this is concerned.
///
/// TODO: move that validation elsewhere? Should probably extract a more general
/// site configuration validator.
fn load_markdown_paths(site_directory: &PathBuf, config: &Config) -> Result<Vec<PathBuf>, String> {
    let site = site_directory
        .to_str()
        .ok_or(format!("bad `site`: {:?}", site_directory));

    let directories = config.directories.content.to_str().ok_or(format!(
        "bad content directory: {:?}",
        config.directories.content
    ));

    let (site, directories) = match (site, directories) {
        (Err(s), Err(d)) => Err(s + "\n" + &d),
        (Err(s), Ok(_)) => Err(s),
        (Ok(_), Err(d)) => Err(d),
        (Ok(s), Ok(d)) => Ok((s, d)),
    }?;

    let content_glob = format!("{}/{}/**/*.md", site, directories);

    let paths = glob(&content_glob).map_err(|err| format!("{:?}", err))?;

    let (valid, errors): (Vec<Result<PathBuf, _>>, Vec<Result<_, String>>) = paths
        .map(|glob_result| glob_result.map_err(|err| format!("{:?}", err)))
        .partition(|path| path.is_ok());

    let errors: Vec<String> = errors.into_iter().map(|e| e.unwrap_err()).collect();
    if errors.len() == 0 {
        Ok(valid.into_iter().map(|v| v.unwrap()).collect())
    } else {
        Err(errors.join(",\n"))
    }
}

// TODO: this is an obvious candidate for a `HashMap`, so use that instead.
type LoadTuple<'a> = (&'a PathBuf, Result<String, String>);

// Instead of using `.collect` at the end of these, I *should* be able to use
// something like `Result<impl Iterator<HashMap<PathBuf, String>, String>`, even
// if the type isn't *quite* taht.
fn load_content<'p>(paths: &'p [PathBuf]) -> Result<Vec<(&'p PathBuf, String)>, String> {
    let (contents, errs): (Vec<LoadTuple>, Vec<LoadTuple>) = paths
        .into_iter()
        .map(|path| {
            let content = std::fs::read_to_string(path).map_err(|e| format!("{:?}", e));
            (path, content)
        })
        .partition(|(_, result)| result.is_ok());

    if errs.len() == 0 {
        Ok(contents
            .into_iter()
            .map(|(path, contents)| (path, contents.unwrap()))
            .collect())
    } else {
        Err(errs
            .into_iter()
            .map(|(_path, err)| err.unwrap_err())
            .collect::<Vec<String>>()
            .join(",\n"))
    }
}

/// Load the templates associated with each taxonomy.
fn load_templates(_site_directory: &PathBuf, _config: &Config) -> Result<Paths, String> {
    unimplemented!()
}

// TODO: put this somewhere else when extracting all the parser logic below into
// its own module.
#[derive(Debug)]
enum ParseState<'s> {
    CodeBlock(&'s syntect::parsing::SyntaxReference),
    PlainTextBlock,
    NonCode,
}

impl<'s> std::default::Default for ParseState<'s> {
    fn default() -> ParseState<'s> {
        ParseState::NonCode
    }
}

enum Syntax {
    Highlight(syntect::highlighting::Theme),
    TagOnly,
    Off,
}

impl Syntax {
    fn try_from(config_syntax: &SyntaxOption) -> Result<Syntax, String> {
        match config_syntax {
            SyntaxOption::Highlight(theme) => {
                let theme_file = PathBuf::from("data").join(theme);
                let theme = ThemeSet::get_theme(theme_file).map_err(|err| format!("{:?}", err))?;
                Ok(Syntax::Highlight(theme))
            }
            SyntaxOption::TagOnly => Ok(Syntax::TagOnly),
            SyntaxOption::Off => Ok(Syntax::Off),
        }
    }
}

/// Generate content from a configuration.
pub fn build(site_directory: PathBuf) -> Result<(), String> {
    // NOTE: this is almost certainly not what we *ultimately* want here, but
    // it'll do as a starting point. There's a lot of `into_iter` and `collect`
    // happening here, which forces us to eagerly materialize a lot of these
    // collections rather than being able to deal with them in a more
    // "streaming" fashion.
    let config = Config::from_file(&PathBuf::from(&site_directory))?;
    let markdown_paths = load_markdown_paths(&site_directory, &config)?;
    let contents = load_content(&markdown_paths)?;

    let configured_syntax = Syntax::try_from(&config.options.syntax)?;

    // TODO: generate these! Best move: on build, generate a `.packdump`, since
    // that seems to be what Syntect knows how to deal with.
    let syntax_dir = PathBuf::from("data/syntaxes");
    let syntax_set =
        syntect::parsing::SyntaxSet::load_from_folder(syntax_dir).expect("can't load syntaxes");

    let parsed_content: Vec<(&PathBuf, String)> = contents
        .into_iter()
        .map(|(path, content)| {
            let parser = Parser::new_ext(&content, pulldown_cmark::Options::all()).scan(
                ParseState::default(),
                |state, event| match event {
                    Event::Start(Tag::CodeBlock(ref language)) => match state {
                        ParseState::NonCode => {
                            *state = match syntax_set.find_syntax_by_name(&language) {
                                Some(syntax) => ParseState::CodeBlock(syntax),
                                None => ParseState::PlainTextBlock,
                            };

                            // We start every code block with `<pre><code class="...">`
                            // so that we always have semantically correct HTML. However, we
                            // don't add the language class if the language isn't set.
                            let content = if language.len() > 0 {
                                Owned(format!(r#"
"#, language))
                            } else {
                                Borrowed("
")
                            };

                            Some(Event::Html(content))
                        }
                        ParseState::CodeBlock(_) | ParseState::PlainTextBlock => {
                            unreachable!("Bad event/state: {:?} with {:?}", event, state);
                        }
                    },

                    Event::End(Tag::CodeBlock(_)) => match state {
                        ParseState::CodeBlock(_) | ParseState::PlainTextBlock => {
                            *state = ParseState::NonCode;
                            Some(Event::Html(Borrowed("
")))
                        }
                        ParseState::NonCode => {
                            unreachable!("Bad event/state: {:?} with {:?}", event, state);
                        }
                    },

                    // When we are in a text block, we *may* be in a code block,
                    // so we may also need to do whatever highlighting is
                    // specified.
                    Event::Text(s) => {
                        match state {
                            ParseState::CodeBlock(syntax) => match &configured_syntax {
                                Syntax::Highlight(theme) => {
                                    let highlighted = syntect::html::highlighted_html_for_string(
                                        &s,
                                        &syntax_set,
                                        &syntax,
                                        theme,
                                    );

                                    Some(Event::Text(Owned(highlighted)))
                                }
                                Syntax::TagOnly => {
                                    // TODO: implement tag-only parsing with syntect
                                    unimplemented!()
                                }
                                Syntax::Off => Some(Event::Text(s)),
                            },
                            ParseState::PlainTextBlock => Some(Event::Text(s)),
                            ParseState::NonCode => Some(Event::Text(s)),
                        }
                    }

                    // TODO: accumulate footnote references and always put them
                    // at the end. Also, generate back-links for them, unless or
                    // until pulldown-cmark does so.
                    Event::FootnoteReference(s) => Some(Event::FootnoteReference(s)),

                    Event::Start(s) => Some(Event::Start(s)),
                    Event::End(e) => Some(Event::End(e)),
                    Event::InlineHtml(s) => Some(Event::InlineHtml(s)),
                    Event::Html(h) => Some(Event::Html(h)),
                    Event::SoftBreak => Some(Event::SoftBreak),
                    Event::HardBreak => Some(Event::HardBreak),
                },
            );

            // These numbers are derived from a brief survey of content in my
            // own website historically, including both code-heavy and
            // prose-heavy posts. It's probably worth exploring other corpora
            // to see if this is basically right.
            let length_bound_estimate = match &config.options.syntax {
                SyntaxOption::Off | SyntaxOption::TagOnly => 2,
                SyntaxOption::Highlight(_) => 8, // TODO: check other pages
            };

            let mut html = String::with_capacity(content.len() * length_bound_estimate);
            cmark::html::push_html(&mut html, parser);

            (path, html)
        })
        .collect();

    // TODO: extract this as part of the writing it out process.
    let output_dir = site_directory.join(&config.directories.output);

    if output_dir.exists() && !output_dir.is_dir() {
        return Err(format!(
            "A file with the same name as the output directory {:?} already exists!",
            output_dir
        ));
    }

    if !output_dir.exists() {
        std::fs::create_dir(&output_dir).expect(&format!(
            "could not create output directory {:?}",
            &output_dir
        ));
    }

    for (path, content) in parsed_content {
        let file_path = path
            .file_name()
            .ok_or(format!("invalid file: {:?}", path))?;

        let dest = output_dir.join(file_path).with_extension("html");

        let mut fd = match File::create(&dest) {
            Ok(file) => file,
            Err(reason) => {
                return Err(format!(
                    "Could not open {} for write: {}",
                    dest.to_string_lossy(),
                    reason
                ));
            }
        };

        let result = write!(fd, "{}", content);
        if let Err(reason) = result {
            return Err(format!("{:?}", reason.kind()));
        }
    }

    Ok(())
}
<|endoftext|>


<commit_before><commit_msg>Update hammer.rs<commit_after><|endoftext|>


<commit_before>extern crate wm_daemons;
use wm_daemons::config::{load_config, load_config_path};
use wm_daemons::dbus_listen::{CallbackMap, SignalInfo, match_signal};

#[macro_use]
extern crate clap;
use clap::{Arg, App};

extern crate config;
use self::config::types::Config;

extern crate dbus;
use self::dbus::{Connection, BusType};

use std::error::Error;
use std::path::Path;

fn run_program(action: &str, conf: &Config) -> () {
    // TODO
}

fn handle_signal(info: &SignalInfo, conf: &Config) -> () {
    if info.member == Some("Lock".to_string()) {
        run_program("lock", conf);
    } else if info.member == Some("Unlock".to_string()) {
        run_program("unlock", conf);
    }
}

fn try_main() -> Result<(), Box<Error>> {
    let matches = App::new("wm-session-agent")
        .version(&crate_version!()[..])
        .author("Ben Boeckel <mathstuf@gmail.com>")
        .about("Listens for logind's Lock and Unlock signals")
        .arg(Arg::with_name("CONFIG")
                .short("c")
                .long("config")
                .help("Path to the configuration value")
                .takes_value(true))
        .arg(Arg::with_name("SESSION")
                .short("s")
                .long("session")
                .required(true)
                .help("Session ID to listen on")
                .takes_value(true))
        .get_matches();

    let conf = try!(if matches.is_present("CONFIG") {
            load_config("wm-session-agent", "config")
        } else {
            load_config_path(Path::new(matches.value_of("CONFIG").unwrap()))
        });

    let sid = matches.value_of("SESSION").unwrap();
    let spath = format!("/org/freedesktop/login1/session/_{}", sid);

    let conn = try!(Connection::get_private(BusType::System));

    let cbs: CallbackMap<Config> = vec![
        (SignalInfo {
            path: None,
            object: None,
            member: Some("Lock".to_string()),
        }, handle_signal),
        (SignalInfo {
            path: None,
            object: None,
            member: Some("Unlock".to_string()),
        }, handle_signal),
    ];

    let match_str = format!("type='signal',interface='org.freedesktop.login1.Session',path='{}'", spath);
    try!(conn.add_match(&match_str[..]));

    for items in conn.iter(100) {
        match_signal(items, &cbs, &conf);
    }

    Ok(())
}

fn main() -> () {
    if let Err(err) = try_main() {
        panic!("{}", err.description());
    }
}
<commit_msg>wm-session-agent: implement the handler<commit_after>extern crate wm_daemons;
use wm_daemons::config::{load_config, load_config_path};
use wm_daemons::dbus_listen::{CallbackMap, SignalInfo, match_signal};
use wm_daemons::exec::run_config_program;

#[macro_use]
extern crate clap;
use clap::{Arg, App};

extern crate config;
use self::config::types::Config;

extern crate dbus;
use self::dbus::{Connection, BusType};

use std::error::Error;
use std::path::Path;

fn run_program(action: &str, conf: &Config) -> () {
    let path = format!("actions.{}", action);
    let res = run_config_program(conf, &path[..]);
    if res.is_err() {
        println!("failed to handle '{}' action: {}", action, res.err().unwrap());
    }
}

fn handle_signal(info: &SignalInfo, conf: &Config) -> () {
    if info.member == Some("Lock".to_string()) {
        run_program("lock", conf);
    } else if info.member == Some("Unlock".to_string()) {
        run_program("unlock", conf);
    }
}

fn try_main() -> Result<(), Box<Error>> {
    let matches = App::new("wm-session-agent")
        .version(&crate_version!()[..])
        .author("Ben Boeckel <mathstuf@gmail.com>")
        .about("Listens for logind's Lock and Unlock signals")
        .arg(Arg::with_name("CONFIG")
                .short("c")
                .long("config")
                .help("Path to the configuration value")
                .takes_value(true))
        .arg(Arg::with_name("SESSION")
                .short("s")
                .long("session")
                .required(true)
                .help("Session ID to listen on")
                .takes_value(true))
        .get_matches();

    let conf = try!(if matches.is_present("CONFIG") {
            load_config("wm-session-agent", "config")
        } else {
            load_config_path(Path::new(matches.value_of("CONFIG").unwrap()))
        });

    let sid = matches.value_of("SESSION").unwrap();
    let spath = format!("/org/freedesktop/login1/session/_{}", sid);

    let conn = try!(Connection::get_private(BusType::System));

    let cbs: CallbackMap<Config> = vec![
        (SignalInfo {
            path: None,
            object: None,
            member: Some("Lock".to_string()),
        }, handle_signal),
        (SignalInfo {
            path: None,
            object: None,
            member: Some("Unlock".to_string()),
        }, handle_signal),
    ];

    let match_str = format!("type='signal',interface='org.freedesktop.login1.Session',path='{}'", spath);
    try!(conn.add_match(&match_str[..]));

    for items in conn.iter(100) {
        match_signal(items, &cbs, &conf);
    }

    Ok(())
}

fn main() -> () {
    if let Err(err) = try_main() {
        panic!("{}", err.description());
    }
}
<|endoftext|>


<commit_before>extern crate semver;

use std::collections::HashMap;
use std::io::net::ip::IpAddr;

#[deriving(PartialEq, Show, Clone)]
pub enum Scheme {
    Http,
    Https
}

#[deriving(PartialEq, Show, Clone)]
pub enum Host<'a> {
    HostName(&'a str),
    HostIp(IpAddr)
}

#[deriving(PartialEq, Show, Clone)]
pub enum Method<'a> {
    Get,
    Post,
    Put,
    Patch,
    Delete,
    Other(&'a str)
}

pub trait Request {
    /// The version of HTTP being used
    fn http_version(&self) -> semver::Version;

    /// The version of the conduit spec being used
    fn conduit_version(&self) -> semver::Version;

    /// The request method, such as GET, POST, PUT, DELETE or PATCH
    fn method<'a>(&'a self) -> Method<'a>;

    /// The scheme part of the request URL
    fn scheme(&self) -> Scheme;

    /// The host part of the requested URL
    fn host<'a>(&'a self) -> Host<'a>;

    /// The initial part of the request URL's path that corresponds
    /// to a virtual root. This allows an application to have a
    /// virtual location that consumes part of the path.
    fn virtual_root<'a>(&'a self) -> Option<&'a str>;

    /// The remainder of the path.
    fn path<'a>(&'a self) -> &'a str;

    /// The portion of the request URL that follows the "?"
    fn query_string<'a>(&'a self) -> Option<&'a str>;

    /// The remote IP address of the client or the last proxy that
    /// sent the request.
    fn remote_ip(&self) -> IpAddr;

    /// The byte-size of the body, if any
    fn content_length(&self) -> Option<uint>;

    /// The request's headers, as conduit::Headers.
    fn headers<'a>(&'a self) -> &'a Headers;

    /// A Reader for the body of the request
    fn body<'a>(&'a mut self) -> &'a mut Reader;
}

pub type HeaderEntries<'a> = &'a Iterator<(&'a str, Vec<&'a str>)>;

pub trait Headers {
    /// Find the value of a given header. Multi-line headers are represented
    /// as an array.
    fn find<'a>(&'a self, key: &str) -> Option<Vec<&'a str>>;

    /// Returns true if a particular header exists
    fn has(&self, key: &str) -> bool;

    /// Iterate over all of the available headers.
    fn iter<'a>(&'a self) -> HeaderEntries<'a>;
}

pub struct Response {
    /// The status code as a tuple of the return code and status string
    pub status: (uint, &'static str),

    /// A Map of the headers
    pub headers: HashMap<String, String>,

    /// A Writer for body of the response
    pub body: Box<Reader + Send>
}

pub type Handler = fn(&mut Request) -> Response;
<commit_msg>Boxed iterator<commit_after>extern crate semver;

use std::collections::HashMap;
use std::io::net::ip::IpAddr;

#[deriving(PartialEq, Show, Clone)]
pub enum Scheme {
    Http,
    Https
}

#[deriving(PartialEq, Show, Clone)]
pub enum Host<'a> {
    HostName(&'a str),
    HostIp(IpAddr)
}

#[deriving(PartialEq, Show, Clone)]
pub enum Method<'a> {
    Get,
    Post,
    Put,
    Patch,
    Delete,
    Other(&'a str)
}

pub trait Request {
    /// The version of HTTP being used
    fn http_version(&self) -> semver::Version;

    /// The version of the conduit spec being used
    fn conduit_version(&self) -> semver::Version;

    /// The request method, such as GET, POST, PUT, DELETE or PATCH
    fn method<'a>(&'a self) -> Method<'a>;

    /// The scheme part of the request URL
    fn scheme(&self) -> Scheme;

    /// The host part of the requested URL
    fn host<'a>(&'a self) -> Host<'a>;

    /// The initial part of the request URL's path that corresponds
    /// to a virtual root. This allows an application to have a
    /// virtual location that consumes part of the path.
    fn virtual_root<'a>(&'a self) -> Option<&'a str>;

    /// The remainder of the path.
    fn path<'a>(&'a self) -> &'a str;

    /// The portion of the request URL that follows the "?"
    fn query_string<'a>(&'a self) -> Option<&'a str>;

    /// The remote IP address of the client or the last proxy that
    /// sent the request.
    fn remote_ip(&self) -> IpAddr;

    /// The byte-size of the body, if any
    fn content_length(&self) -> Option<uint>;

    /// The request's headers, as conduit::Headers.
    fn headers<'a>(&'a self) -> &'a Headers;

    /// A Reader for the body of the request
    fn body<'a>(&'a mut self) -> &'a mut Reader;
}

pub type HeaderEntries<'a> = Box<Iterator<(&'a str, Vec<&'a str>)>>;

pub trait Headers {
    /// Find the value of a given header. Multi-line headers are represented
    /// as an array.
    fn find<'a>(&'a self, key: &str) -> Option<Vec<&'a str>>;

    /// Returns true if a particular header exists
    fn has(&self, key: &str) -> bool;

    /// Iterate over all of the available headers.
    fn iter<'a>(&'a self) -> HeaderEntries<'a>;
}

pub struct Response {
    /// The status code as a tuple of the return code and status string
    pub status: (uint, &'static str),

    /// A Map of the headers
    pub headers: HashMap<String, String>,

    /// A Writer for body of the response
    pub body: Box<Reader + Send>
}

pub type Handler = fn(&mut Request) -> Response;
<|endoftext|>