<commit_before>//! In this example, you will be shown various ways of sharing data between events and commands.
//! And how to use locks correctly to avoid deadlocking the bot.

use std::{sync::{Arc, atomic::{AtomicUsize, Ordering}}, collections::HashMap, env};

use serenity::{
    async_trait,
    model::{channel::Message, gateway::Ready},
    prelude::*,

    framework::standard::{
        Args, CommandResult, StandardFramework,
        macros::{command, group, hook},
    },
};

use tokio::sync::RwLock;

// A container type is created for inserting into the Client's `data`, which
// allows for data to be accessible across all events and framework commands, or
// anywhere else that has a copy of the `data` Arc.
// These places are usually where either Context or Client is present.
//
// Documentation about TypeMap can be found here:
// https://docs.rs/typemap_rev/0.1/typemap_rev/struct.TypeMap.html
struct CommandCounter;

impl TypeMapKey for CommandCounter {
    type Value = Arc<RwLock<HashMap<String, u64>>>;
}

struct MessageCount;

impl TypeMapKey for MessageCount {
    // While you will be using RwLock or Mutex most of the time you want to modify data,
    // sometimes it's not required; like for example, with static data, or if you are using other
    // kinds of atomic operators.
    //
    // Arc should stay, to allow for the data lock to be closed early.
    type Value = Arc<AtomicUsize>;
}

#[group]
#[commands(ping, command_usage, owo_count)]
struct General;

#[hook]
async fn before(ctx: &Context, msg: &Message, command_name: &str) -> bool {
    println!("Running command '{}' invoked by '{}'", command_name, msg.author.tag());

    let counter_lock = {
        // While data is a RwLock, it's recommended that you always open the lock as read.
        // This is mainly done to avoid Deadlocks for having a possible writer waiting for multiple
        // readers to close.
        let data_read = ctx.data.read().await;

        // Since the CommandCounter Value is wrapped in an Arc, cloning will not duplicate the
        // data, instead the reference is cloned.
        // We wap every value on in an Arc, as to keep the data lock open for the least time possible,
        // to again, avoid deadlocking it.
        data_read.get::<CommandCounter>().expect("Expected CommandCounter in TypeMap.").clone()
    };

    // Just like with client.data in main, we want to keep write locks open the least time
    // possible, so we wrap them on a block so they get automatically closed at the end.
    {
        // The HashMap of CommandCounter is wrapped in an RwLock; since we want to write to it, we will
        // open the lock in write mode.
        let mut counter = counter_lock.write().await;

        // And we write the amount of times the command has been called to it.
        let entry = counter.entry(command_name.to_string()).or_insert(0);
        *entry += 1;
    }

    true
}

struct Handler;

#[async_trait]
impl EventHandler for Handler {
    async fn message(&self, ctx: Context, msg: Message) {
        if msg.content.to_lowercase().contains("owo") {
            // Since data is located in Context, this means you are also able to use it within events!
            let count = {
                let data_read = ctx.data.read().await;
                data_read.get::<MessageCount>().expect("Expected MessageCount in TypeMap.").clone()
            };

            // Atomic operations with ordering do not require mut to be modified.
            // In this case, we want to increase the message count by 1.
            // https://doc.rust-lang.org/std/sync/atomic/struct.AtomicUsize.html#method.fetch_add
            count.fetch_add(1, Ordering::SeqCst);
        }
    }

    async fn ready(&self, _: Context, ready: Ready) {
        println!("{} is connected!", ready.user.name);
    }
}

#[tokio::main]
async fn main() {
    let token = env::var("DISCORD_TOKEN").expect("Expected a token in the environment");

    let framework = StandardFramework::new()
        .configure(|c| c
            .with_whitespace(true)
            .prefix("~")
        )
        .before(before)
        .group(&GENERAL_GROUP);

    let mut client = Client::builder(&token)
        .event_handler(Handler)
        .framework(framework)
        .await
        .expect("Err creating client");

    // This is wherere we can initially insert the data we desire into the "global" data TypeMap.
    // client.data is wrapped on a RwLock, and since we want to insert to it, we have to open it in
    // write mode, but there's a small thing catch:
    // There can only be a single writer to a given lock open in the entire application, this means
    // you can't open a new write lock until the previous write lock has closed.
    // This is not the case with read locks, read locks can be open indefinitely, BUT as soon as
    // you need to open the lock in write mode, all the read locks must be closed.
    //
    // You can find more information about deadlocks in the Rust Book, ch16-03:
    // https://doc.rust-lang.org/book/ch16-03-shared-state.html
    //
    // All of this means that we have to keep locks open for the least time possible, so we put
    // them inside a block, so they get closed automatically when droped.
    // If we don't do this, we would never be able to open the data lock anywhere else.
    {
        // Open the data lock in write mode, so keys can be inserted to it.
        let mut data = client.data.write().await;

        // The CommandCounter Value has the following type:
        // Arc<RwLock<HashMap<String, u64>>>
        // So, we have to insert the same type to it.
        data.insert::<CommandCounter>(Arc::new(RwLock::new(HashMap::default())));

        data.insert::<MessageCount>(Arc::new(AtomicUsize::new(0)));
    }

    if let Err(why) = client.start().await {
        eprintln!("Client error: {:?}", why);
    }
}

#[command]
async fn ping(ctx: &Context, msg: &Message) -> CommandResult {
    msg.reply(ctx, "Pong!").await?;

    Ok(())
}

/// Usage: `~command_usage <command_name>`
/// Example: `~command_usage ping`
#[command]
async fn command_usage(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
    let command_name = match args.single_quoted::<String>() {
        Ok(x) => x,
        Err(_) => {
            msg.reply(ctx, "I require an argument to run this command.").await?;
            return Ok(());
        }
    };

    // Yet again, we want to keep the locks open for the least time possible.
    let amount = {
        // Since we only want to read the data and not write to it, we open it in read mode,
        // and since this is open in read mode, it means that there can be multiple locks open at
        // the same time, and as mentioned earlier, it's heavily recommended that you only open
        // the data lock in read mode, as it will avoid a lot of possible deadlocks.
        let data_read = ctx.data.read().await;

        // Then we obtain the value we need from data, in this case, we want the command counter.
        // The returned value from get() is an Arc, so the reference will be cloned, rather than
        // the data.
        let command_counter_lock = data_read.get::<CommandCounter>().expect("Expected CommandCounter in TypeMap.").clone();

        let command_counter = command_counter_lock.read().await;
        // And we return a usable value from it.
        // This time, the value is not Arc, so the data will be cloned.
        command_counter.get(&command_name).map_or(0, |x| *x)
    };

    if amount == 0 {
        msg.reply(ctx, format!("The command `{}` has not yet been used.", command_name)).await?;
    } else {
        msg.reply(ctx, format!("The command `{}` has been used {} time/s this session!", command_name, amount)).await?;
    }

    Ok(())
}

#[command]
async fn owo_count(ctx: &Context, msg: &Message) -> CommandResult {
    let raw_count = {
        let data_read = ctx.data.read().await;
        data_read.get::<MessageCount>().expect("Expected MessageCount in TypeMap.").clone()
    };

    let count = raw_count.load(Ordering::Relaxed);

    if count == 1 {
        msg.reply(ctx, "You are the first one to say owo this session! *because it's on the command name* :P").await?;
    } else {
        msg.reply(ctx, format!("OWO Has been said {} times!", count)).await?;
    }

    Ok(())
}
<commit_msg>Fix a small typo in example 14 (#1030)<commit_after>//! In this example, you will be shown various ways of sharing data between events and commands.
//! And how to use locks correctly to avoid deadlocking the bot.

use std::{sync::{Arc, atomic::{AtomicUsize, Ordering}}, collections::HashMap, env};

use serenity::{
    async_trait,
    model::{channel::Message, gateway::Ready},
    prelude::*,

    framework::standard::{
        Args, CommandResult, StandardFramework,
        macros::{command, group, hook},
    },
};

use tokio::sync::RwLock;

// A container type is created for inserting into the Client's `data`, which
// allows for data to be accessible across all events and framework commands, or
// anywhere else that has a copy of the `data` Arc.
// These places are usually where either Context or Client is present.
//
// Documentation about TypeMap can be found here:
// https://docs.rs/typemap_rev/0.1/typemap_rev/struct.TypeMap.html
struct CommandCounter;

impl TypeMapKey for CommandCounter {
    type Value = Arc<RwLock<HashMap<String, u64>>>;
}

struct MessageCount;

impl TypeMapKey for MessageCount {
    // While you will be using RwLock or Mutex most of the time you want to modify data,
    // sometimes it's not required; like for example, with static data, or if you are using other
    // kinds of atomic operators.
    //
    // Arc should stay, to allow for the data lock to be closed early.
    type Value = Arc<AtomicUsize>;
}

#[group]
#[commands(ping, command_usage, owo_count)]
struct General;

#[hook]
async fn before(ctx: &Context, msg: &Message, command_name: &str) -> bool {
    println!("Running command '{}' invoked by '{}'", command_name, msg.author.tag());

    let counter_lock = {
        // While data is a RwLock, it's recommended that you always open the lock as read.
        // This is mainly done to avoid Deadlocks for having a possible writer waiting for multiple
        // readers to close.
        let data_read = ctx.data.read().await;

        // Since the CommandCounter Value is wrapped in an Arc, cloning will not duplicate the
        // data, instead the reference is cloned.
        // We wap every value on in an Arc, as to keep the data lock open for the least time possible,
        // to again, avoid deadlocking it.
        data_read.get::<CommandCounter>().expect("Expected CommandCounter in TypeMap.").clone()
    };

    // Just like with client.data in main, we want to keep write locks open the least time
    // possible, so we wrap them on a block so they get automatically closed at the end.
    {
        // The HashMap of CommandCounter is wrapped in an RwLock; since we want to write to it, we will
        // open the lock in write mode.
        let mut counter = counter_lock.write().await;

        // And we write the amount of times the command has been called to it.
        let entry = counter.entry(command_name.to_string()).or_insert(0);
        *entry += 1;
    }

    true
}

struct Handler;

#[async_trait]
impl EventHandler for Handler {
    async fn message(&self, ctx: Context, msg: Message) {
        if msg.content.to_lowercase().contains("owo") {
            // Since data is located in Context, this means you are also able to use it within events!
            let count = {
                let data_read = ctx.data.read().await;
                data_read.get::<MessageCount>().expect("Expected MessageCount in TypeMap.").clone()
            };

            // Atomic operations with ordering do not require mut to be modified.
            // In this case, we want to increase the message count by 1.
            // https://doc.rust-lang.org/std/sync/atomic/struct.AtomicUsize.html#method.fetch_add
            count.fetch_add(1, Ordering::SeqCst);
        }
    }

    async fn ready(&self, _: Context, ready: Ready) {
        println!("{} is connected!", ready.user.name);
    }
}

#[tokio::main]
async fn main() {
    let token = env::var("DISCORD_TOKEN").expect("Expected a token in the environment");

    let framework = StandardFramework::new()
        .configure(|c| c
            .with_whitespace(true)
            .prefix("~")
        )
        .before(before)
        .group(&GENERAL_GROUP);

    let mut client = Client::builder(&token)
        .event_handler(Handler)
        .framework(framework)
        .await
        .expect("Err creating client");

    // This is where we can initially insert the data we desire into the "global" data TypeMap.
    // client.data is wrapped on a RwLock, and since we want to insert to it, we have to open it in
    // write mode, but there's a small thing catch:
    // There can only be a single writer to a given lock open in the entire application, this means
    // you can't open a new write lock until the previous write lock has closed.
    // This is not the case with read locks, read locks can be open indefinitely, BUT as soon as
    // you need to open the lock in write mode, all the read locks must be closed.
    //
    // You can find more information about deadlocks in the Rust Book, ch16-03:
    // https://doc.rust-lang.org/book/ch16-03-shared-state.html
    //
    // All of this means that we have to keep locks open for the least time possible, so we put
    // them inside a block, so they get closed automatically when droped.
    // If we don't do this, we would never be able to open the data lock anywhere else.
    {
        // Open the data lock in write mode, so keys can be inserted to it.
        let mut data = client.data.write().await;

        // The CommandCounter Value has the following type:
        // Arc<RwLock<HashMap<String, u64>>>
        // So, we have to insert the same type to it.
        data.insert::<CommandCounter>(Arc::new(RwLock::new(HashMap::default())));

        data.insert::<MessageCount>(Arc::new(AtomicUsize::new(0)));
    }

    if let Err(why) = client.start().await {
        eprintln!("Client error: {:?}", why);
    }
}

#[command]
async fn ping(ctx: &Context, msg: &Message) -> CommandResult {
    msg.reply(ctx, "Pong!").await?;

    Ok(())
}

/// Usage: `~command_usage <command_name>`
/// Example: `~command_usage ping`
#[command]
async fn command_usage(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
    let command_name = match args.single_quoted::<String>() {
        Ok(x) => x,
        Err(_) => {
            msg.reply(ctx, "I require an argument to run this command.").await?;
            return Ok(());
        }
    };

    // Yet again, we want to keep the locks open for the least time possible.
    let amount = {
        // Since we only want to read the data and not write to it, we open it in read mode,
        // and since this is open in read mode, it means that there can be multiple locks open at
        // the same time, and as mentioned earlier, it's heavily recommended that you only open
        // the data lock in read mode, as it will avoid a lot of possible deadlocks.
        let data_read = ctx.data.read().await;

        // Then we obtain the value we need from data, in this case, we want the command counter.
        // The returned value from get() is an Arc, so the reference will be cloned, rather than
        // the data.
        let command_counter_lock = data_read.get::<CommandCounter>().expect("Expected CommandCounter in TypeMap.").clone();

        let command_counter = command_counter_lock.read().await;
        // And we return a usable value from it.
        // This time, the value is not Arc, so the data will be cloned.
        command_counter.get(&command_name).map_or(0, |x| *x)
    };

    if amount == 0 {
        msg.reply(ctx, format!("The command `{}` has not yet been used.", command_name)).await?;
    } else {
        msg.reply(ctx, format!("The command `{}` has been used {} time/s this session!", command_name, amount)).await?;
    }

    Ok(())
}

#[command]
async fn owo_count(ctx: &Context, msg: &Message) -> CommandResult {
    let raw_count = {
        let data_read = ctx.data.read().await;
        data_read.get::<MessageCount>().expect("Expected MessageCount in TypeMap.").clone()
    };

    let count = raw_count.load(Ordering::Relaxed);

    if count == 1 {
        msg.reply(ctx, "You are the first one to say owo this session! *because it's on the command name* :P").await?;
    } else {
        msg.reply(ctx, format!("OWO Has been said {} times!", count)).await?;
    }

    Ok(())
}
<|endoftext|>


<commit_before>//! Conversions between ULID and UUID.

use crate::Ulid;
use uuid::Uuid;

impl From<Uuid> for Ulid {
    fn from(uuid: Uuid) -> Self {
        Ulid(uuid.as_u128())
    }
}

impl From<&Ulid> for Uuid {
    fn from(ulid: &Ulid) -> Self {
        Uuid::from_u128(ulid.0)
    }
}
<commit_msg>add uuid tests<commit_after>//! Conversions between ULID and UUID.

use crate::Ulid;
use uuid::Uuid;

impl From<Uuid> for Ulid {
    fn from(uuid: Uuid) -> Self {
        Ulid(uuid.as_u128())
    }
}

impl From<Ulid> for Uuid {
    fn from(ulid: Ulid) -> Self {
        Uuid::from_u128(ulid.0)
    }
}

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

    #[test]
    fn uuid_cycle() {
        let ulid = Ulid::new();
        let uuid: Uuid = ulid.into();
        let ulid2: Ulid = uuid.into();

        assert_eq!(ulid, ulid2);
    }

    #[test]
    fn uuid_str_cycle() {
        let uuid_txt = "771a3bce-02e9-4428-a68e-b1e7e82b7f9f";
        let ulid_txt = "3Q38XWW0Q98GMAD3NHWZM2PZWZ";

        let ulid: Ulid = Uuid::parse_str(uuid_txt).unwrap().into();
        assert_eq!(ulid.to_string(), ulid_txt);

        let uuid: Uuid = ulid.into();
        assert_eq!(uuid.to_string(), uuid_txt);
    }
}
<|endoftext|>


<commit_before>use std::fmt;
use std::marker::PhantomData;
use std::str::FromStr;
use std::sync::Arc;
use std::thread;
use std::time::{Duration, Instant};

use antidote::RwLock;
use irc::client::prelude::*;
use rlua::prelude::*;
use rlua::HookTriggers;

use chrono::NaiveDateTime;
use time;

use crate::plugin::*;
use crate::FrippyClient;
pub mod database;
use self::database::Database;

mod utils;
use self::utils::*;
use crate::utils::Url;

use self::error::*;
use crate::error::ErrorKind as FrippyErrorKind;
use crate::error::FrippyError;
use failure::{format_err, ResultExt};

use frippy_derive::PluginName;

static LUA_SANDBOX: &'static str = include_str!("sandbox.lua");

#[derive(PluginName)]
pub struct Factoid<T: Database, C: Client> {
    factoids: RwLock<T>,
    phantom: PhantomData<C>,
}

impl<T: Database, C: Client> Factoid<T, C> {
    pub fn new(db: T) -> Self {
        Factoid {
            factoids: RwLock::new(db),
            phantom: PhantomData,
        }
    }

    fn create_factoid(
        &self,
        name: &str,
        content: &str,
        author: &str,
    ) -> Result<&str, FactoidError> {
        let count = self.factoids.read().count_factoids(name)?;
        let tm = time::now().to_timespec();

        let factoid = database::NewFactoid {
            name,
            idx: count,
            content,
            author,
            created: NaiveDateTime::from_timestamp(tm.sec, 0u32),
        };

        Ok(self
            .factoids
            .write()
            .insert_factoid(&factoid)
            .map(|()| "Successfully added!")?)
    }

    fn add(&self, command: &mut PluginCommand) -> Result<&str, FactoidError> {
        if command.tokens.len() < 2 {
            Err(ErrorKind::InvalidCommand)?;
        }

        let name = command.tokens.remove(0);
        let content = command.tokens.join(" ");

        Ok(self.create_factoid(&name, &content, &command.source)?)
    }

    fn add_from_url(&self, command: &mut PluginCommand) -> Result<&str, FactoidError> {
        if command.tokens.len() < 2 {
            Err(ErrorKind::InvalidCommand)?;
        }

        let name = command.tokens.remove(0);
        let url = &command.tokens[0];
        let content = Url::from(url.as_ref())
            .max_kib(1024)
            .request()
            .context(ErrorKind::Download)?;

        Ok(self.create_factoid(&name, &content, &command.source)?)
    }

    fn remove(&self, command: &mut PluginCommand) -> Result<&str, FactoidError> {
        if command.tokens.is_empty() {
            Err(ErrorKind::InvalidCommand)?;
        }

        let name = command.tokens.remove(0);
        let count = self.factoids.read().count_factoids(&name)?;

        match self.factoids.write().delete_factoid(&name, count - 1) {
            Ok(()) => Ok("Successfully removed"),
            Err(e) => Err(e)?,
        }
    }

    fn get(&self, command: &PluginCommand) -> Result<String, FactoidError> {
        let (name, idx) = match command.tokens.len() {
            0 => Err(ErrorKind::InvalidCommand)?,
            1 => {
                let name = &command.tokens[0];
                let count = self.factoids.read().count_factoids(name)?;

                if count < 1 {
                    Err(ErrorKind::NotFound)?;
                }

                (name, count - 1)
            }
            _ => {
                let name = &command.tokens[0];
                let idx = match i32::from_str(&command.tokens[1]) {
                    Ok(i) => i,
                    Err(_) => Err(ErrorKind::InvalidCommand)?,
                };

                (name, idx)
            }
        };

        let factoid = self
            .factoids
            .read()
            .get_factoid(name, idx)
            .context(ErrorKind::NotFound)?;

        let mut message = factoid.content.replace("\n", "|").replace("\r", "");
        message.truncate(512);

        Ok(format!("{}: {}", factoid.name, message))
    }

    fn info(&self, command: &PluginCommand) -> Result<String, FactoidError> {
        match command.tokens.len() {
            0 => Err(ErrorKind::InvalidCommand)?,
            1 => {
                let name = &command.tokens[0];
                let count = self.factoids.read().count_factoids(name)?;

                Ok(match count {
                    0 => Err(ErrorKind::NotFound)?,
                    1 => format!("There is 1 version of {}", name),
                    _ => format!("There are {} versions of {}", count, name),
                })
            }
            _ => {
                let name = &command.tokens[0];
                let idx = i32::from_str(&command.tokens[1]).context(ErrorKind::InvalidIndex)?;
                let factoid = self.factoids.read().get_factoid(name, idx)?;

                Ok(format!(
                    "{}: Added by {} at {} UTC",
                    name, factoid.author, factoid.created
                ))
            }
        }
    }

    fn exec(&self, mut command: PluginCommand) -> Result<String, FactoidError> {
        if command.tokens.is_empty() {
            Err(ErrorKind::InvalidIndex)?
        } else {
            let name = command.tokens.remove(0);
            let count = self.factoids.read().count_factoids(&name)?;
            let factoid = self.factoids.read().get_factoid(&name, count - 1)?;

            let content = factoid.content;
            let mut message = if content.starts_with('>') {
                let content = String::from(&content[1..]);

                if content.starts_with('>') {
                    content
                } else {
                    match self.run_lua(&name, &content, &command) {
                        Ok(v) => v,
                        Err(e) => match e {
                            LuaError::CallbackError { cause, .. } => match *cause {
                                LuaError::MemoryError(_) => {
                                    String::from("memory error: Factoid used over 1 MiB of ram")
                                }
                                _ => cause.to_string(),
                            },
                            LuaError::MemoryError(_) => {
                                String::from("memory error: Factoid used over 1 MiB of ram")
                            }
                            _ => e.to_string(),
                        },
                    }
                }
            } else {
                content
            };

            message.truncate(512);
            Ok(message.replace("\n", "|").replace("\r", ""))
        }
    }

    fn run_lua(&self, name: &str, code: &str, command: &PluginCommand) -> Result<String, LuaError> {
        let args = command
            .tokens
            .iter()
            .filter(|x| !x.is_empty())
            .map(ToOwned::to_owned)
            .collect::<Vec<String>>();

        let lua = Lua::new();
        // TODO Is this actually 1 Mib?
        lua.set_memory_limit(Some(1024 * 1024));

        let start = Instant::now();
        // Check if the factoid timed out
        lua.set_hook(
            HookTriggers {
                every_line: true,
                ..Default::default()
            },
            move |_, _| {
                if Instant::now() - start > Duration::from_secs(30) {
                    return Err(LuaError::ExternalError(Arc::new(
                        format_err!("Factoid timed out after 30 seconds").compat(),
                    )));
                }

                // Limit the cpu usage of factoids
                thread::sleep(Duration::from_millis(1));

                Ok(())
            },
        );

        let output = lua.context(|ctx| {
            let globals = ctx.globals();

            globals.set("factoid", code)?;
            globals.set(
                "download",
                ctx.create_function(|ctx, url| download(&ctx, url))?,
            )?;
            globals.set(
                "json_decode",
                ctx.create_function(|ctx, json| json_decode(&ctx, json))?,
            )?;
            globals.set("sleep", ctx.create_function(|ctx, ms| sleep(&ctx, ms))?)?;
            globals.set("args", args)?;
            globals.set("input", command.tokens.join(" "))?;
            globals.set("user", command.source.clone())?;
            globals.set("channel", command.target.clone())?;
            globals.set("output", ctx.create_table()?)?;

            ctx.load(LUA_SANDBOX).set_name(name)?.exec()?;

            Ok(globals.get::<_, Vec<String>>("output")?)
        })?;

        Ok(output.join("|"))
    }

    fn help(&self) -> &str {
        "usage: factoids <subcommand>\r\n\
         subcommands: add, fromurl, remove, get, info, exec, help"
    }
}

impl<T: Database, C: FrippyClient> Plugin for Factoid<T, C> {
    type Client = C;
    fn execute(&self, _: &Self::Client, message: &Message) -> ExecutionStatus {
        match message.command {
            Command::PRIVMSG(_, ref content) => {
                if content.starts_with('!') {
                    ExecutionStatus::RequiresThread
                } else {
                    ExecutionStatus::Done
                }
            }
            _ => ExecutionStatus::Done,
        }
    }

    fn execute_threaded(
        &self,
        client: &Self::Client,
        message: &Message,
    ) -> Result<(), FrippyError> {
        if let Command::PRIVMSG(_, mut content) = message.command.clone() {
            content.remove(0);

            let t: Vec<String> = content.split(' ').map(ToOwned::to_owned).collect();

            let c = PluginCommand {
                source: message.source_nickname().unwrap().to_owned(),
                target: message.response_target().unwrap().to_owned(),
                tokens: t,
            };

            if let Ok(f) = self.exec(c) {
                client
                    .send_privmsg(&message.response_target().unwrap(), &f)
                    .context(FrippyErrorKind::Connection)?;
            }
        }

        Ok(())
    }

    fn command(
        &self,
        client: &Self::Client,
        mut command: PluginCommand,
    ) -> Result<(), FrippyError> {
        if command.tokens.is_empty() {
            client
                .send_privmsg(&command.target, "Invalid command")
                .context(FrippyErrorKind::Connection)?;

            return Ok(());
        }

        let target = command.target.clone();

        let sub_command = command.tokens.remove(0);
        let result = match sub_command.as_ref() {
            "add" => self.add(&mut command).map(|s| s.to_owned()),
            "fromurl" => self.add_from_url(&mut command).map(|s| s.to_owned()),
            "remove" => self.remove(&mut command).map(|s| s.to_owned()),
            "get" => self.get(&command),
            "info" => self.info(&command),
            "exec" => self.exec(command),
            "help" => Ok(self.help().to_owned()),
            _ => Err(ErrorKind::InvalidCommand.into()),
        };

        match result {
            Ok(m) => {
                client
                    .send_privmsg(&target, &m)
                    .context(FrippyErrorKind::Connection)?;
            }
            Err(e) => {
                let message = e.to_string();
                client
                    .send_privmsg(&target, &message)
                    .context(FrippyErrorKind::Connection)?;
                Err(e).context(FrippyErrorKind::Factoid)?
            }
        }

        Ok(())
    }

    fn evaluate(&self, _: &Self::Client, _: PluginCommand) -> Result<String, String> {
        Err(String::from(
            "Evaluation of commands is not implemented for Factoid at this time",
        ))
    }
}

impl<T: Database, C: FrippyClient> fmt::Debug for Factoid<T, C> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Factoid {{ ... }}")
    }
}

pub mod error {
    use failure::Fail;
    use frippy_derive::Error;

    #[derive(Copy, Clone, Eq, PartialEq, Debug, Fail, Error)]
    #[error = "FactoidError"]
    pub enum ErrorKind {
        /// Invalid command error
        #[fail(display = "Invalid Command")]
        InvalidCommand,

        /// Invalid index error
        #[fail(display = "Invalid index")]
        InvalidIndex,

        /// Download  error
        #[fail(display = "Download failed")]
        Download,

        /// Duplicate error
        #[fail(display = "Entry already exists")]
        Duplicate,

        /// Not found error
        #[fail(display = "Factoid was not found")]
        NotFound,

        /// MySQL error
        #[cfg(feature = "mysql")]
        #[fail(display = "Failed to execute MySQL Query")]
        MysqlError,

        /// No connection error
        #[cfg(feature = "mysql")]
        #[fail(display = "No connection to the database")]
        NoConnection,
    }
}
<commit_msg>Leave 100 bytes when truncating factoids<commit_after>use std::fmt;
use std::marker::PhantomData;
use std::str::FromStr;
use std::sync::Arc;
use std::thread;
use std::time::{Duration, Instant};

use antidote::RwLock;
use irc::client::prelude::*;
use rlua::prelude::*;
use rlua::HookTriggers;

use chrono::NaiveDateTime;
use time;

use crate::plugin::*;
use crate::FrippyClient;
pub mod database;
use self::database::Database;

mod utils;
use self::utils::*;
use crate::utils::Url;

use self::error::*;
use crate::error::ErrorKind as FrippyErrorKind;
use crate::error::FrippyError;
use failure::{format_err, ResultExt};

use frippy_derive::PluginName;

static LUA_SANDBOX: &'static str = include_str!("sandbox.lua");

#[derive(PluginName)]
pub struct Factoid<T: Database, C: Client> {
    factoids: RwLock<T>,
    phantom: PhantomData<C>,
}

impl<T: Database, C: Client> Factoid<T, C> {
    pub fn new(db: T) -> Self {
        Factoid {
            factoids: RwLock::new(db),
            phantom: PhantomData,
        }
    }

    fn create_factoid(
        &self,
        name: &str,
        content: &str,
        author: &str,
    ) -> Result<&str, FactoidError> {
        let count = self.factoids.read().count_factoids(name)?;
        let tm = time::now().to_timespec();

        let factoid = database::NewFactoid {
            name,
            idx: count,
            content,
            author,
            created: NaiveDateTime::from_timestamp(tm.sec, 0u32),
        };

        Ok(self
            .factoids
            .write()
            .insert_factoid(&factoid)
            .map(|()| "Successfully added!")?)
    }

    fn add(&self, command: &mut PluginCommand) -> Result<&str, FactoidError> {
        if command.tokens.len() < 2 {
            Err(ErrorKind::InvalidCommand)?;
        }

        let name = command.tokens.remove(0);
        let content = command.tokens.join(" ");

        Ok(self.create_factoid(&name, &content, &command.source)?)
    }

    fn add_from_url(&self, command: &mut PluginCommand) -> Result<&str, FactoidError> {
        if command.tokens.len() < 2 {
            Err(ErrorKind::InvalidCommand)?;
        }

        let name = command.tokens.remove(0);
        let url = &command.tokens[0];
        let content = Url::from(url.as_ref())
            .max_kib(1024)
            .request()
            .context(ErrorKind::Download)?;

        Ok(self.create_factoid(&name, &content, &command.source)?)
    }

    fn remove(&self, command: &mut PluginCommand) -> Result<&str, FactoidError> {
        if command.tokens.is_empty() {
            Err(ErrorKind::InvalidCommand)?;
        }

        let name = command.tokens.remove(0);
        let count = self.factoids.read().count_factoids(&name)?;

        match self.factoids.write().delete_factoid(&name, count - 1) {
            Ok(()) => Ok("Successfully removed"),
            Err(e) => Err(e)?,
        }
    }

    fn get(&self, command: &PluginCommand) -> Result<String, FactoidError> {
        let (name, idx) = match command.tokens.len() {
            0 => Err(ErrorKind::InvalidCommand)?,
            1 => {
                let name = &command.tokens[0];
                let count = self.factoids.read().count_factoids(name)?;

                if count < 1 {
                    Err(ErrorKind::NotFound)?;
                }

                (name, count - 1)
            }
            _ => {
                let name = &command.tokens[0];
                let idx = match i32::from_str(&command.tokens[1]) {
                    Ok(i) => i,
                    Err(_) => Err(ErrorKind::InvalidCommand)?,
                };

                (name, idx)
            }
        };

        let factoid = self
            .factoids
            .read()
            .get_factoid(name, idx)
            .context(ErrorKind::NotFound)?;

        let mut message = factoid.content.replace("\n", "|").replace("\r", "");
        message.truncate(512);

        Ok(format!("{}: {}", factoid.name, message))
    }

    fn info(&self, command: &PluginCommand) -> Result<String, FactoidError> {
        match command.tokens.len() {
            0 => Err(ErrorKind::InvalidCommand)?,
            1 => {
                let name = &command.tokens[0];
                let count = self.factoids.read().count_factoids(name)?;

                Ok(match count {
                    0 => Err(ErrorKind::NotFound)?,
                    1 => format!("There is 1 version of {}", name),
                    _ => format!("There are {} versions of {}", count, name),
                })
            }
            _ => {
                let name = &command.tokens[0];
                let idx = i32::from_str(&command.tokens[1]).context(ErrorKind::InvalidIndex)?;
                let factoid = self.factoids.read().get_factoid(name, idx)?;

                Ok(format!(
                    "{}: Added by {} at {} UTC",
                    name, factoid.author, factoid.created
                ))
            }
        }
    }

    fn exec(&self, mut command: PluginCommand) -> Result<String, FactoidError> {
        if command.tokens.is_empty() {
            Err(ErrorKind::InvalidIndex)?
        } else {
            let name = command.tokens.remove(0);
            let count = self.factoids.read().count_factoids(&name)?;
            let factoid = self.factoids.read().get_factoid(&name, count - 1)?;

            let content = factoid.content;
            let mut message = if content.starts_with('>') {
                let content = String::from(&content[1..]);

                if content.starts_with('>') {
                    content
                } else {
                    match self.run_lua(&name, &content, &command) {
                        Ok(v) => v,
                        Err(e) => match e {
                            LuaError::CallbackError { cause, .. } => match *cause {
                                LuaError::MemoryError(_) => {
                                    String::from("memory error: Factoid used over 1 MiB of ram")
                                }
                                _ => cause.to_string(),
                            },
                            LuaError::MemoryError(_) => {
                                String::from("memory error: Factoid used over 1 MiB of ram")
                            }
                            _ => e.to_string(),
                        },
                    }
                }
            } else {
                content
            };

            message.truncate(412);
            Ok(message.replace("\n", "|").replace("\r", ""))
        }
    }

    fn run_lua(&self, name: &str, code: &str, command: &PluginCommand) -> Result<String, LuaError> {
        let args = command
            .tokens
            .iter()
            .filter(|x| !x.is_empty())
            .map(ToOwned::to_owned)
            .collect::<Vec<String>>();

        let lua = Lua::new();
        // TODO Is this actually 1 Mib?
        lua.set_memory_limit(Some(1024 * 1024));

        let start = Instant::now();
        // Check if the factoid timed out
        lua.set_hook(
            HookTriggers {
                every_line: true,
                ..Default::default()
            },
            move |_, _| {
                if Instant::now() - start > Duration::from_secs(30) {
                    return Err(LuaError::ExternalError(Arc::new(
                        format_err!("Factoid timed out after 30 seconds").compat(),
                    )));
                }

                // Limit the cpu usage of factoids
                thread::sleep(Duration::from_millis(1));

                Ok(())
            },
        );

        let output = lua.context(|ctx| {
            let globals = ctx.globals();

            globals.set("factoid", code)?;
            globals.set(
                "download",
                ctx.create_function(|ctx, url| download(&ctx, url))?,
            )?;
            globals.set(
                "json_decode",
                ctx.create_function(|ctx, json| json_decode(&ctx, json))?,
            )?;
            globals.set("sleep", ctx.create_function(|ctx, ms| sleep(&ctx, ms))?)?;
            globals.set("args", args)?;
            globals.set("input", command.tokens.join(" "))?;
            globals.set("user", command.source.clone())?;
            globals.set("channel", command.target.clone())?;
            globals.set("output", ctx.create_table()?)?;

            ctx.load(LUA_SANDBOX).set_name(name)?.exec()?;

            Ok(globals.get::<_, Vec<String>>("output")?)
        })?;

        Ok(output.join("|"))
    }

    fn help(&self) -> &str {
        "usage: factoids <subcommand>\r\n\
         subcommands: add, fromurl, remove, get, info, exec, help"
    }
}

impl<T: Database, C: FrippyClient> Plugin for Factoid<T, C> {
    type Client = C;
    fn execute(&self, _: &Self::Client, message: &Message) -> ExecutionStatus {
        match message.command {
            Command::PRIVMSG(_, ref content) => {
                if content.starts_with('!') {
                    ExecutionStatus::RequiresThread
                } else {
                    ExecutionStatus::Done
                }
            }
            _ => ExecutionStatus::Done,
        }
    }

    fn execute_threaded(
        &self,
        client: &Self::Client,
        message: &Message,
    ) -> Result<(), FrippyError> {
        if let Command::PRIVMSG(_, mut content) = message.command.clone() {
            content.remove(0);

            let t: Vec<String> = content.split(' ').map(ToOwned::to_owned).collect();

            let c = PluginCommand {
                source: message.source_nickname().unwrap().to_owned(),
                target: message.response_target().unwrap().to_owned(),
                tokens: t,
            };

            if let Ok(f) = self.exec(c) {
                client
                    .send_privmsg(&message.response_target().unwrap(), &f)
                    .context(FrippyErrorKind::Connection)?;
            }
        }

        Ok(())
    }

    fn command(
        &self,
        client: &Self::Client,
        mut command: PluginCommand,
    ) -> Result<(), FrippyError> {
        if command.tokens.is_empty() {
            client
                .send_privmsg(&command.target, "Invalid command")
                .context(FrippyErrorKind::Connection)?;

            return Ok(());
        }

        let target = command.target.clone();

        let sub_command = command.tokens.remove(0);
        let result = match sub_command.as_ref() {
            "add" => self.add(&mut command).map(|s| s.to_owned()),
            "fromurl" => self.add_from_url(&mut command).map(|s| s.to_owned()),
            "remove" => self.remove(&mut command).map(|s| s.to_owned()),
            "get" => self.get(&command),
            "info" => self.info(&command),
            "exec" => self.exec(command),
            "help" => Ok(self.help().to_owned()),
            _ => Err(ErrorKind::InvalidCommand.into()),
        };

        match result {
            Ok(m) => {
                client
                    .send_privmsg(&target, &m)
                    .context(FrippyErrorKind::Connection)?;
            }
            Err(e) => {
                let message = e.to_string();
                client
                    .send_privmsg(&target, &message)
                    .context(FrippyErrorKind::Connection)?;
                Err(e).context(FrippyErrorKind::Factoid)?
            }
        }

        Ok(())
    }

    fn evaluate(&self, _: &Self::Client, _: PluginCommand) -> Result<String, String> {
        Err(String::from(
            "Evaluation of commands is not implemented for Factoid at this time",
        ))
    }
}

impl<T: Database, C: FrippyClient> fmt::Debug for Factoid<T, C> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Factoid {{ ... }}")
    }
}

pub mod error {
    use failure::Fail;
    use frippy_derive::Error;

    #[derive(Copy, Clone, Eq, PartialEq, Debug, Fail, Error)]
    #[error = "FactoidError"]
    pub enum ErrorKind {
        /// Invalid command error
        #[fail(display = "Invalid Command")]
        InvalidCommand,

        /// Invalid index error
        #[fail(display = "Invalid index")]
        InvalidIndex,

        /// Download  error
        #[fail(display = "Download failed")]
        Download,

        /// Duplicate error
        #[fail(display = "Entry already exists")]
        Duplicate,

        /// Not found error
        #[fail(display = "Factoid was not found")]
        NotFound,

        /// MySQL error
        #[cfg(feature = "mysql")]
        #[fail(display = "Failed to execute MySQL Query")]
        MysqlError,

        /// No connection error
        #[cfg(feature = "mysql")]
        #[fail(display = "No connection to the database")]
        NoConnection,
    }
}
<|endoftext|>


<commit_before>//! A set of builders to generate Rust source for PHF data structures at
//! compile time.
//!
//! The provided builders are intended to be used in a Cargo build script to
//! generate a Rust source file that will be included in a library at build
//! time.
//!
//! # Examples
//!
//! build.rs
//!
//! ```rust,no_run
//! extern crate phf_codegen;
//!
//! use std::env;
//! use std::fs::File;
//! use std::io::{BufWriter, Write};
//! use std::path::Path;
//!
//! fn main() {
//!     let path = Path::new(&env::var("OUT_DIR").unwrap()).join("codegen.rs");
//!     let mut file = BufWriter::new(File::create(&path).unwrap());
//!
//! write!(&mut file, "static KEYWORDS: phf::Map<&'static str, Keyword> =
//! ").unwrap();
//!     phf_codegen::Map::new()
//!         .entry("loop", "Keyword::Loop")
//!         .entry("continue", "Keyword::Continue")
//!         .entry("break", "Keyword::Break")
//!         .entry("fn", "Keyword::Fn")
//!         .entry("extern", "Keyword::Extern")
//!         .build(&mut file)
//!         .unwrap();
//!     write!(&mut file, ";\n").unwrap();
//! }
//! ```
//!
//! lib.rs
//!
//! ```ignore
//! extern crate phf;
//!
//! #[derive(Clone)]
//! enum Keyword {
//!     Loop,
//!     Continue,
//!     Break,
//!     Fn,
//!     Extern,
//! }
//!
//! include!(concat!(env!("OUT_DIR"), "/codegen.rs"));
//!
//! pub fn parse_keyword(keyword: &str) -> Option<Keyword> {
//!     KEYWORDS.get(keyword).cloned()
//! }
//! ```
//!
//! # Note
//!
//! The compiler's stack will overflow when processing extremely long method
//! chains (500+ calls). When generating large PHF data structures, consider
//! looping over the entries or making each call a separate statement:
//!
//! ```rust
//! let entries = [("hello", "1"), ("world", "2")];
//!
//! let mut builder = phf_codegen::Map::new();
//! for &(key, value) in &entries {
//!     builder.entry(key, value);
//! }
//! // ...
//! ```
//!
//! ```rust
//! let mut builder = phf_codegen::Map::new();
//! builder.entry("hello", "1");
//! builder.entry("world", "2");
//! // ...
//! ```
#![doc(html_root_url="http://sfackler.github.io/rust-phf/doc/v0.7.14")]
extern crate phf_shared;
extern crate phf_generator;

use phf_shared::PhfHash;
use std::collections::HashSet;
use std::fmt;
use std::hash::Hash;
use std::io;
use std::io::prelude::*;

/// A builder for the `phf::Map` type.
pub struct Map<K> {
    keys: Vec<K>,
    values: Vec<String>,
}

impl<K: Hash+PhfHash+Eq+fmt::Debug> Map<K> {
    /// Creates a new `phf::Map` builder.
    pub fn new() -> Map<K> {
        // FIXME rust#27438
        //
        // On Windows/MSVC there are major problems with the handling of dllimport.
        // Here, because downstream build scripts only invoke generics from phf_codegen,
        // the linker ends up throwing a way a bunch of static symbols we actually need.
        // This works around the problem, assuming that all clients call `Map::new` by
        // calling a non-generic function.
        fn noop_fix_for_27438() {
        }
        noop_fix_for_27438();

        Map {
            keys: vec![],
            values: vec![],
        }
    }

    /// Adds an entry to the builder.
    ///
    /// `value` will be written exactly as provided in the constructed source.
    pub fn entry(&mut self, key: K, value: &str) -> &mut Map<K> {
        self.keys.push(key);
        self.values.push(value.to_string());
        self
    }

    /// Constructs a `phf::Map`, outputting Rust source to the provided writer.
    ///
    /// # Panics
    ///
    /// Panics if there are any duplicate keys.
    pub fn build<W: Write>(&self, w: &mut W) -> io::Result<()> {
        let mut set = HashSet::new();
        for key in &self.keys {
            if !set.insert(key) {
                panic!("duplicate key `{:?}`", key);
            }
        }

        let state = phf_generator::generate_hash(&self.keys);

        try!(write!(w,
                    "::phf::Map {{
    key: {},
    disps: ::phf::Slice::Static(&[",
                    state.key));
        for &(d1, d2) in &state.disps {
            try!(write!(w,
                        "
        ({}, {}),",
                        d1,
                        d2));
        }
        try!(write!(w,
                    "
    ]),
    entries: ::phf::Slice::Static(&["));
        for &idx in &state.map {
            try!(write!(w,
                        "
        ({:?}, {}),",
                        &self.keys[idx],
                        &self.values[idx]));
        }
        write!(w,
               "
    ]),
}}")
    }
}

/// A builder for the `phf::Set` type.
pub struct Set<T> {
    map: Map<T>,
}

impl<T: Hash+PhfHash+Eq+fmt::Debug> Set<T> {
    /// Constructs a new `phf::Set` builder.
    pub fn new() -> Set<T> {
        Set { map: Map::new() }
    }

    /// Adds an entry to the builder.
    pub fn entry(&mut self, entry: T) -> &mut Set<T> {
        self.map.entry(entry, "()");
        self
    }

    /// Constructs a `phf::Set`, outputting Rust source to the provided writer.
    ///
    /// # Panics
    ///
    /// Panics if there are any duplicate entries.
    pub fn build<W: Write>(&self, w: &mut W) -> io::Result<()> {
        try!(write!(w, "::phf::Set {{ map: "));
        try!(self.map.build(w));
        write!(w, " }}")
    }
}

/// A builder for the `phf::OrderedMap` type.
pub struct OrderedMap<K> {
    keys: Vec<K>,
    values: Vec<String>,
}

impl<K: Hash+PhfHash+Eq+fmt::Debug> OrderedMap<K> {
    /// Constructs a enw `phf::OrderedMap` builder.
    pub fn new() -> OrderedMap<K> {
        OrderedMap {
            keys: vec![],
            values: vec![],
        }
    }

    /// Adds an entry to the builder.
    ///
    /// `value` will be written exactly as provided in the constructed source.
    pub fn entry(&mut self, key: K, value: &str) -> &mut OrderedMap<K> {
        self.keys.push(key);
        self.values.push(value.to_string());
        self
    }

    /// Constructs a `phf::OrderedMap`, outputting Rust source to the provided
    /// writer.
    ///
    /// # Panics
    ///
    /// Panics if there are any duplicate keys.
    pub fn build<W: Write>(&self, w: &mut W) -> io::Result<()> {
        let mut set = HashSet::new();
        for key in &self.keys {
            if !set.insert(key) {
                panic!("duplicate key `{:?}`", key);
            }
        }

        let state = phf_generator::generate_hash(&self.keys);

        try!(write!(w,
                    "::phf::OrderedMap {{
    key: {},
    disps: ::phf::Slice::Static(&[",
                    state.key));
        for &(d1, d2) in &state.disps {
            try!(write!(w,
                        "
        ({}, {}),",
                        d1,
                        d2));
        }
        try!(write!(w,
                    "
    ]),
    idxs: ::phf::Slice::Static(&["));
        for &idx in &state.map {
            try!(write!(w,
                        "
        {},",
                        idx));
        }
        try!(write!(w,
                    "
    ]),
    entries: ::phf::Slice::Static(&["));
        for (key, value) in self.keys.iter().zip(self.values.iter()) {
            try!(write!(w,
                        "
        ({:?}, {}),",
                        key,
                        value));
        }
        write!(w,
               "
    ]),
}}")
    }
}

/// A builder for the `phf::OrderedSet` type.
pub struct OrderedSet<T> {
    map: OrderedMap<T>,
}

impl<T: Hash+PhfHash+Eq+fmt::Debug> OrderedSet<T> {
    /// Constructs a new `phf::OrderedSet` builder.
    pub fn new() -> OrderedSet<T> {
        OrderedSet { map: OrderedMap::new() }
    }

    /// Adds an entry to the builder.
    pub fn entry(&mut self, entry: T) -> &mut OrderedSet<T> {
        self.map.entry(entry, "()");
        self
    }

    /// Constructs a `phf::OrderedSet`, outputting Rust source to the provided
    /// writer.
    ///
    /// # Panics
    ///
    /// Panics if there are any duplicate entries.
    pub fn build<W: Write>(&self, w: &mut W) -> io::Result<()> {
        try!(write!(w, "::phf::OrderedSet {{ map: "));
        try!(self.map.build(w));
        write!(w, " }}")
    }
}
<commit_msg>Allow specifying a module path to the phf crate for codegen<commit_after>//! A set of builders to generate Rust source for PHF data structures at
//! compile time.
//!
//! The provided builders are intended to be used in a Cargo build script to
//! generate a Rust source file that will be included in a library at build
//! time.
//!
//! # Examples
//!
//! build.rs
//!
//! ```rust,no_run
//! extern crate phf_codegen;
//!
//! use std::env;
//! use std::fs::File;
//! use std::io::{BufWriter, Write};
//! use std::path::Path;
//!
//! fn main() {
//!     let path = Path::new(&env::var("OUT_DIR").unwrap()).join("codegen.rs");
//!     let mut file = BufWriter::new(File::create(&path).unwrap());
//!
//! write!(&mut file, "static KEYWORDS: phf::Map<&'static str, Keyword> =
//! ").unwrap();
//!     phf_codegen::Map::new()
//!         .entry("loop", "Keyword::Loop")
//!         .entry("continue", "Keyword::Continue")
//!         .entry("break", "Keyword::Break")
//!         .entry("fn", "Keyword::Fn")
//!         .entry("extern", "Keyword::Extern")
//!         .build(&mut file)
//!         .unwrap();
//!     write!(&mut file, ";\n").unwrap();
//! }
//! ```
//!
//! lib.rs
//!
//! ```ignore
//! extern crate phf;
//!
//! #[derive(Clone)]
//! enum Keyword {
//!     Loop,
//!     Continue,
//!     Break,
//!     Fn,
//!     Extern,
//! }
//!
//! include!(concat!(env!("OUT_DIR"), "/codegen.rs"));
//!
//! pub fn parse_keyword(keyword: &str) -> Option<Keyword> {
//!     KEYWORDS.get(keyword).cloned()
//! }
//! ```
//!
//! # Note
//!
//! The compiler's stack will overflow when processing extremely long method
//! chains (500+ calls). When generating large PHF data structures, consider
//! looping over the entries or making each call a separate statement:
//!
//! ```rust
//! let entries = [("hello", "1"), ("world", "2")];
//!
//! let mut builder = phf_codegen::Map::new();
//! for &(key, value) in &entries {
//!     builder.entry(key, value);
//! }
//! // ...
//! ```
//!
//! ```rust
//! let mut builder = phf_codegen::Map::new();
//! builder.entry("hello", "1");
//! builder.entry("world", "2");
//! // ...
//! ```
#![doc(html_root_url="http://sfackler.github.io/rust-phf/doc/v0.7.14")]
extern crate phf_shared;
extern crate phf_generator;

use phf_shared::PhfHash;
use std::collections::HashSet;
use std::fmt;
use std::hash::Hash;
use std::io;
use std::io::prelude::*;

/// A builder for the `phf::Map` type.
pub struct Map<K> {
    keys: Vec<K>,
    values: Vec<String>,
    path: String,
}

impl<K: Hash+PhfHash+Eq+fmt::Debug> Map<K> {
    /// Creates a new `phf::Map` builder.
    pub fn new() -> Map<K> {
        // FIXME rust#27438
        //
        // On Windows/MSVC there are major problems with the handling of dllimport.
        // Here, because downstream build scripts only invoke generics from phf_codegen,
        // the linker ends up throwing a way a bunch of static symbols we actually need.
        // This works around the problem, assuming that all clients call `Map::new` by
        // calling a non-generic function.
        fn noop_fix_for_27438() {
        }
        noop_fix_for_27438();

        Map {
            keys: vec![],
            values: vec![],
            path: String::from("::phf"),
        }
    }

    /// Set the path to the `phf` crate from the global namespace
    pub fn phf_path(&mut self, path: &str) -> &mut Map<K> {
        self.path = path.to_owned();
        self
    }

    /// Adds an entry to the builder.
    ///
    /// `value` will be written exactly as provided in the constructed source.
    pub fn entry(&mut self, key: K, value: &str) -> &mut Map<K> {
        self.keys.push(key);
        self.values.push(value.to_owned());
        self
    }

    /// Constructs a `phf::Map`, outputting Rust source to the provided writer.
    ///
    /// # Panics
    ///
    /// Panics if there are any duplicate keys.
    pub fn build<W: Write>(&self, w: &mut W) -> io::Result<()> {
        let mut set = HashSet::new();
        for key in &self.keys {
            if !set.insert(key) {
                panic!("duplicate key `{:?}`", key);
            }
        }

        let state = phf_generator::generate_hash(&self.keys);

        try!(write!(w,
                    "{}::Map {{
    key: {},
    disps: {}::Slice::Static(&[",
                    self.path, state.key, self.path));
        for &(d1, d2) in &state.disps {
            try!(write!(w,
                        "
        ({}, {}),",
                        d1,
                        d2));
        }
        try!(write!(w,
                    "
    ]),
    entries: ::phf::Slice::Static(&["));
        for &idx in &state.map {
            try!(write!(w,
                        "
        ({:?}, {}),",
                        &self.keys[idx],
                        &self.values[idx]));
        }
        write!(w,
               "
    ]),
}}")
    }
}

/// A builder for the `phf::Set` type.
pub struct Set<T> {
    map: Map<T>,
}

impl<T: Hash+PhfHash+Eq+fmt::Debug> Set<T> {
    /// Constructs a new `phf::Set` builder.
    pub fn new() -> Set<T> {
        Set {
            map: Map::new(),
        }
    }

    /// Set the path to the `phf` crate from the global namespace
    pub fn phf_path(&mut self, path: &str) -> &mut Set<T> {
        self.map.phf_path(path);
        self
    }

    /// Adds an entry to the builder.
    pub fn entry(&mut self, entry: T) -> &mut Set<T> {
        self.map.entry(entry, "()");
        self
    }

    /// Constructs a `phf::Set`, outputting Rust source to the provided writer.
    ///
    /// # Panics
    ///
    /// Panics if there are any duplicate entries.
    pub fn build<W: Write>(&self, w: &mut W) -> io::Result<()> {
        try!(write!(w, "{}::Set {{ map: ", self.map.path));
        try!(self.map.build(w));
        write!(w, " }}")
    }
}

/// A builder for the `phf::OrderedMap` type.
pub struct OrderedMap<K> {
    keys: Vec<K>,
    values: Vec<String>,
    path: String,
}

impl<K: Hash+PhfHash+Eq+fmt::Debug> OrderedMap<K> {
    /// Constructs a enw `phf::OrderedMap` builder.
    pub fn new() -> OrderedMap<K> {
        OrderedMap {
            keys: vec![],
            values: vec![],
            path: String::from("::phf"),
        }
    }

    /// Set the path to the `phf` crate from the global namespace
    pub fn phf_path(&mut self, path: &str) -> &mut OrderedMap<K> {
        self.path = path.to_owned();
        self
    }

    /// Adds an entry to the builder.
    ///
    /// `value` will be written exactly as provided in the constructed source.
    pub fn entry(&mut self, key: K, value: &str) -> &mut OrderedMap<K> {
        self.keys.push(key);
        self.values.push(value.to_owned());
        self
    }

    /// Constructs a `phf::OrderedMap`, outputting Rust source to the provided
    /// writer.
    ///
    /// # Panics
    ///
    /// Panics if there are any duplicate keys.
    pub fn build<W: Write>(&self, w: &mut W) -> io::Result<()> {
        let mut set = HashSet::new();
        for key in &self.keys {
            if !set.insert(key) {
                panic!("duplicate key `{:?}`", key);
            }
        }

        let state = phf_generator::generate_hash(&self.keys);

        try!(write!(w,
                    "{}::OrderedMap {{
    key: {},
    disps: {}::Slice::Static(&[",
                    self.path, state.key, self.path));
        for &(d1, d2) in &state.disps {
            try!(write!(w,
                        "
        ({}, {}),",
                        d1,
                        d2));
        }
        try!(write!(w,
                    "
    ]),
    idxs: ::phf::Slice::Static(&["));
        for &idx in &state.map {
            try!(write!(w,
                        "
        {},",
                        idx));
        }
        try!(write!(w,
                    "
    ]),
    entries: ::phf::Slice::Static(&["));
        for (key, value) in self.keys.iter().zip(self.values.iter()) {
            try!(write!(w,
                        "
        ({:?}, {}),",
                        key,
                        value));
        }
        write!(w,
               "
    ]),
}}")
    }
}

/// A builder for the `phf::OrderedSet` type.
pub struct OrderedSet<T> {
    map: OrderedMap<T>,
}

impl<T: Hash+PhfHash+Eq+fmt::Debug> OrderedSet<T> {
    /// Constructs a new `phf::OrderedSet` builder.
    pub fn new() -> OrderedSet<T> {
        OrderedSet {
            map: OrderedMap::new(),
        }
    }

    /// Set the path to the `phf` crate from the global namespace
    pub fn phf_path(&mut self, path: &str) -> &mut OrderedSet<T> {
        self.map.phf_path(path);
        self
    }

    /// Adds an entry to the builder.
    pub fn entry(&mut self, entry: T) -> &mut OrderedSet<T> {
        self.map.entry(entry, "()");
        self
    }

    /// Constructs a `phf::OrderedSet`, outputting Rust source to the provided
    /// writer.
    ///
    /// # Panics
    ///
    /// Panics if there are any duplicate entries.
    pub fn build<W: Write>(&self, w: &mut W) -> io::Result<()> {
        try!(write!(w, "{}::OrderedSet {{ map: ", self.map.path));
        try!(self.map.build(w));
        write!(w, " }}")
    }
}
<|endoftext|>


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

//! `Transaction` related types.

use std::borrow::Cow;

use messages::Message;
use storage::{Fork, StorageValue};
use crypto::Hash;
use encoding::serialize::json::ExonumJson;

// TODO: Remove attribute when `const fn` becomes stable.
#[cfg_attr(feature="cargo-clippy", allow(cast_lossless))]
static MAX_RESERVED_VALUE: u16 = ::std::u8::MAX as u16;

/// A trait that describes transaction processing rules (a group of sequential operations
/// with the Exonum storage) for the given `Message`.
pub trait Transaction: Message + ExonumJson + 'static {
    /// Verifies the transaction, which includes the message signature verification and other
    /// specific internal constraints. verify is intended to check the internal consistency of
    /// a transaction; it has no access to the blockchain state.
    /// If a transaction fails verify, it is considered incorrect and cannot be included into
    /// any correct block proposal. Incorrect transactions are never included into the blockchain.
    ///
    /// *This method should not use external data, that is, it must be a pure function.*
    ///
    /// # Examples
    ///
    /// ```
    /// use exonum::blockchain::Transaction;
    /// use exonum::crypto::PublicKey;
    /// # use exonum::blockchain::TransactionStatus;
    /// # use exonum::storage::Fork;
    ///
    /// message! {
    ///     struct MyTransaction {
    ///         // Transaction definition...
    /// #       const TYPE = 1;
    /// #       const ID = 1;
    /// #       const SIZE = 32;
    ///         field public_key:    &PublicKey    [0 => 32]
    ///     }
    /// }
    ///
    /// impl Transaction for MyTransaction {
    ///     fn verify(&self) -> bool {
    ///         self.verify_signature(self.public_key())
    ///     }
    ///
    ///     // Other methods...
    ///     // ...
    /// #   fn execute(&self, _: &mut Fork) -> TransactionStatus { TransactionStatus::Succeeded }
    /// }
    fn verify(&self) -> bool;

    /// Takes the current blockchain state via `fork` and can modify it if certain conditions
    /// are met.
    ///
    /// # Notes
    ///
    /// - When programming `execute`, you should perform state-related checks before any changes
    /// to the state and return early if these checks fail.
    /// - If the execute method of a transaction raises a `panic`, the changes made by the
    /// transactions are discarded, but the transaction itself is still considered committed.
    /// - A transaction execution status (see `ExecutionStatus` for the details) is stored in the
    /// blockchain and can be accessed through api.
    ///
    ///
    /// # Examples
    ///
    /// ```
    /// use exonum::blockchain::Transaction;
    /// use exonum::crypto::PublicKey;
    /// use exonum::storage::Fork;
    ///
    /// message! {
    ///     struct MyTransaction {
    ///         // Transaction definition...
    /// #       const TYPE = 1;
    /// #       const ID = 1;
    /// #       const SIZE = 32;
    /// #       field public_key:    &PublicKey    [0 => 32]
    ///     }
    /// }
    ///
    /// impl Transaction for MyTransaction {
    ///     fn execute(&self, fork: &mut Fork) {
    ///         // Read and/or write into storage.
    ///         // ...
    ///
    ///         // Returns execution status.
    ///         TransactionStatus::Succeeded
    ///     }
    ///
    ///     // Other methods...
    ///     // ...
    /// #   fn verify(&self) -> bool { true }
    /// }
    fn execute(&self, fork: &mut Fork) -> TransactionStatus;
}

/// Execution status of the transaction.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum TransactionStatus {
    /// Successful transaction execution.
    Succeeded,
    /// Panic occurred during transaction execution. This status should not be used explicitly.
    Panic,
    /// General failure (unspecified reason).
    Failed,
    /// User defined execution status. Can have different meanings for different transactions and
    /// services.
    Custom(u8),
}

impl<'a, T: Transaction> From<T> for Box<Transaction + 'a> {
    fn from(tx: T) -> Self {
        Box::new(tx) as Box<Transaction>
    }
}

impl StorageValue for TransactionStatus {
    fn hash(&self) -> Hash {
        u16::hash(&from_status(*self))
    }

    fn into_bytes(self) -> Vec<u8> {
        u16::into_bytes(from_status(self))
    }

    fn from_bytes(value: Cow<[u8]>) -> Self {
        to_status(u16::from_bytes(value))
    }
}

fn from_status(status: TransactionStatus) -> u16 {
    match status {
        TransactionStatus::Succeeded => 0,
        TransactionStatus::Panic => 1,
        TransactionStatus::Failed => 2,
        TransactionStatus::Custom(value) => u16::from(value) + MAX_RESERVED_VALUE + 1,
    }
}

fn to_status(value: u16) -> TransactionStatus {
    match value {
        0 => TransactionStatus::Succeeded,
        1 => TransactionStatus::Panic,
        2 => TransactionStatus::Failed,
        val => {
            assert!(
                val > MAX_RESERVED_VALUE,
                "Invalid ExecutionStatus value: {}",
                val
            );
            TransactionStatus::Custom((val - MAX_RESERVED_VALUE - 1) as u8)
        }
    }
}

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

    #[test]
    fn execution_status_round_trip() {
        let statuses = [
            TransactionStatus::Succeeded,
            TransactionStatus::Panic,
            TransactionStatus::Failed,
            TransactionStatus::Custom(0),
            TransactionStatus::Custom(1),
            TransactionStatus::Custom(100),
            TransactionStatus::Custom(255),
        ];

        for status in &statuses {
            let bytes = status.clone().into_bytes();
            let new_status = TransactionStatus::from_bytes(Cow::Borrowed(&bytes));

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

//! `Transaction` related types.

use std::borrow::Cow;

use messages::Message;
use storage::{Fork, StorageValue};
use crypto::Hash;
use encoding::serialize::json::ExonumJson;

// TODO: Remove attribute when `const fn` becomes stable.
#[cfg_attr(feature="cargo-clippy", allow(cast_lossless))]
static MAX_RESERVED_VALUE: u16 = ::std::u8::MAX as u16;

/// A trait that describes transaction processing rules (a group of sequential operations
/// with the Exonum storage) for the given `Message`.
pub trait Transaction: Message + ExonumJson + 'static {
    /// Verifies the transaction, which includes the message signature verification and other
    /// specific internal constraints. verify is intended to check the internal consistency of
    /// a transaction; it has no access to the blockchain state.
    /// If a transaction fails verify, it is considered incorrect and cannot be included into
    /// any correct block proposal. Incorrect transactions are never included into the blockchain.
    ///
    /// *This method should not use external data, that is, it must be a pure function.*
    ///
    /// # Examples
    ///
    /// ```
    /// # #[macro_use] extern crate exonum;
    /// #
    /// use exonum::blockchain::Transaction;
    /// use exonum::crypto::PublicKey;
    /// use exonum::messages::Message;
    /// # use exonum::blockchain::TransactionStatus;
    /// # use exonum::storage::Fork;
    ///
    /// message! {
    ///     struct MyTransaction {
    ///         // Transaction definition...
    /// #       const TYPE = 1;
    /// #       const ID = 1;
    /// #       const SIZE = 32;
    ///         field public_key:    &PublicKey    [0 => 32]
    ///     }
    /// }
    ///
    /// impl Transaction for MyTransaction {
    ///     fn verify(&self) -> bool {
    ///         self.verify_signature(self.public_key())
    ///     }
    ///
    ///     // Other methods...
    ///     // ...
    /// #   fn execute(&self, _: &mut Fork) -> TransactionStatus { TransactionStatus::Succeeded }
    /// }
    /// # fn main() {}
    fn verify(&self) -> bool;

    /// Takes the current blockchain state via `fork` and can modify it if certain conditions
    /// are met.
    ///
    /// # Notes
    ///
    /// - When programming `execute`, you should perform state-related checks before any changes
    /// to the state and return early if these checks fail.
    /// - If the execute method of a transaction raises a `panic`, the changes made by the
    /// transactions are discarded, but the transaction itself is still considered committed.
    /// - A transaction execution status (see `ExecutionStatus` for the details) is stored in the
    /// blockchain and can be accessed through api.
    ///
    ///
    /// # Examples
    ///
    /// ```
    /// # #[macro_use] extern crate exonum;
    /// #
    /// use exonum::blockchain::{Transaction, TransactionStatus};
    /// use exonum::crypto::PublicKey;
    /// use exonum::storage::Fork;
    ///
    /// message! {
    ///     struct MyTransaction {
    ///         // Transaction definition...
    /// #       const TYPE = 1;
    /// #       const ID = 1;
    /// #       const SIZE = 32;
    /// #       field public_key:    &PublicKey    [0 => 32]
    ///     }
    /// }
    ///
    /// impl Transaction for MyTransaction {
    ///     fn execute(&self, fork: &mut Fork) -> TransactionStatus {
    ///         // Read and/or write into storage.
    ///         // ...
    ///
    ///         // Returns execution status.
    ///         TransactionStatus::Succeeded
    ///     }
    ///
    ///     // Other methods...
    ///     // ...
    /// #   fn verify(&self) -> bool { true }
    /// }
    /// # fn main() {}
    fn execute(&self, fork: &mut Fork) -> TransactionStatus;
}

/// Execution status of the transaction.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum TransactionStatus {
    /// Successful transaction execution.
    Succeeded,
    /// Panic occurred during transaction execution. This status should not be used explicitly.
    Panic,
    /// General failure (unspecified reason).
    Failed,
    /// User defined execution status. Can have different meanings for different transactions and
    /// services.
    Custom(u8),
}

impl<'a, T: Transaction> From<T> for Box<Transaction + 'a> {
    fn from(tx: T) -> Self {
        Box::new(tx) as Box<Transaction>
    }
}

impl StorageValue for TransactionStatus {
    fn hash(&self) -> Hash {
        u16::hash(&from_status(*self))
    }

    fn into_bytes(self) -> Vec<u8> {
        u16::into_bytes(from_status(self))
    }

    fn from_bytes(value: Cow<[u8]>) -> Self {
        to_status(u16::from_bytes(value))
    }
}

fn from_status(status: TransactionStatus) -> u16 {
    match status {
        TransactionStatus::Succeeded => 0,
        TransactionStatus::Panic => 1,
        TransactionStatus::Failed => 2,
        TransactionStatus::Custom(value) => u16::from(value) + MAX_RESERVED_VALUE + 1,
    }
}

fn to_status(value: u16) -> TransactionStatus {
    match value {
        0 => TransactionStatus::Succeeded,
        1 => TransactionStatus::Panic,
        2 => TransactionStatus::Failed,
        val => {
            assert!(
                val > MAX_RESERVED_VALUE,
                "Invalid ExecutionStatus value: {}",
                val
            );
            TransactionStatus::Custom((val - MAX_RESERVED_VALUE - 1) as u8)
        }
    }
}

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

    #[test]
    fn execution_status_round_trip() {
        let statuses = [
            TransactionStatus::Succeeded,
            TransactionStatus::Panic,
            TransactionStatus::Failed,
            TransactionStatus::Custom(0),
            TransactionStatus::Custom(1),
            TransactionStatus::Custom(100),
            TransactionStatus::Custom(255),
        ];

        for status in &statuses {
            let bytes = status.clone().into_bytes();
            let new_status = TransactionStatus::from_bytes(Cow::Borrowed(&bytes));

            assert_eq!(*status, new_status);
        }
    }
}
<|endoftext|>


<commit_before>use nom::{IResult, is_space};
use std::collections::{HashMap, HashSet};
use std::collections::hash_map::Entry;
use std::fmt::Debug;
use std::hash::Hash;
use std::iter::FromIterator;
use std::str::FromStr;

use util::parsing::*;

/// Structure containing the elements of type `A` in a equivalence class of type `B`
#[derive(Debug, Eq, PartialEq)]
pub struct EquivalenceClass<A, B>
    where A: Eq + Hash,
{
    label: B,
    set: Option<HashSet<A>>,
}

impl<A, B> From<(B, Option<Vec<A>>)> for EquivalenceClass<A, B>
    where A: Eq + Hash,
{
    fn from((label, set): (B, Option<Vec<A>>)) -> EquivalenceClass<A, B> {
        EquivalenceClass {
            label,
            set: set.and_then(|elements| Some(HashSet::from_iter(elements))),
        }
    }
}

/// A struct containing a remapping of elements of type `A` into their respective equivalence classes of type `B`.
#[derive(Clone, Debug)]
pub struct EquivalenceRelation<A, B>
    where A: Eq + Hash,
{
    map: HashMap<A, B>,
    default: B,
}

impl<A, B> PartialEq for EquivalenceRelation<A, B>
    where A: Eq + Hash,
          B: PartialEq,
{
    fn eq(&self, other: &Self) -> bool {
        self.default == other.default && self.map == other.map
    }
}

impl<A, B> EquivalenceRelation<A, B>
    where A: Eq + Hash,
          B: Clone + Eq + Hash,
{
    pub fn new(map: HashMap<B, HashSet<A>>, default: B) -> Self {
        match Self::new_safe(map, default) {
            Ok(relation) => relation,
            Err(error) => panic!(error),
        }
    }

    pub fn new_safe(map: HashMap<B, HashSet<A>>, default: B) -> Result<Self, String> {
        let mut relation_map = HashMap::new();
        for (class_name, members) in map {
            if class_name == default {
                return Err(String::from(
                    "There can only be one default class in the equivalence relation!"
                ));
            }

            for value in members {
                match relation_map.entry(value) {
                    Entry::Vacant(v) => {
                        v.insert(class_name.clone());
                    },
                    Entry::Occupied(_) => {
                        return Err(String::from(
                            "All classes of the equivalence relation must be disjoint!"
                        ));
                    },
                }
            }
        }

        Ok(EquivalenceRelation { map: relation_map, default })
    }

    // returns the equivalence class of a given value
    pub fn project(&self, key: &A) -> B {
        match self.map.get(key) {
            None => self.default.clone(),
            Some(e) => e.clone(),
        }
    }
}

impl<A, B> From<Vec<EquivalenceClass<A, B>>> for EquivalenceRelation<A, B>
    where A: Eq + Hash,
          B: Clone + Eq + Hash,
{
    fn from(classes: Vec<EquivalenceClass<A, B>>) -> EquivalenceRelation<A, B> {
        let mut map = HashMap::new();
        let mut default = None;

        for EquivalenceClass { label, set } in classes {
            if let Some(elements) = set {
                for element in elements {
                    map.insert(element, label.clone());
                }
            } else {
                default = Some(label);
            }
        }

        EquivalenceRelation { map, default: default.unwrap() }
    }
}

impl<A, B> FromStr for EquivalenceRelation<A, B>
    where A: Clone + Eq + Hash + FromStr,
          A::Err: Debug,
          B: Clone + Eq + Hash + FromStr,
          B::Err: Debug,
{
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let mut map = HashMap::new();
        let mut default = match parse_token(s.as_bytes()) {
            IResult::Done(_, result) => result,
            _                        => return Err(format!("Could not parse {}", s)),
        };
        for l in s.lines() {
            if !l.is_empty() {
                match l.trim().parse()? {
                    EquivalenceClass { label, set: Some(elements) } => {
                        map.insert(label, elements);
                    },
                    EquivalenceClass { label, set: None } => {
                        default = label;
                    },
                }
            }
        }

        Ok(EquivalenceRelation::new(map, default))

    }
}

impl <A, B> FromStr for EquivalenceClass<A, B>
    where A: Eq + FromStr + Hash,
          A::Err: Debug,
          B: FromStr,
          B::Err: Debug,
{
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match parse_set(s.as_bytes()) {
            IResult::Done(_, result) => Ok(result),
            _                        => Err(format!("Could not parse {}", s))
        }
    }
}

fn parse_set<A, B>(input: &[u8]) -> IResult<&[u8], EquivalenceClass<A, B>>
    where A: Eq + FromStr + Hash,
          A::Err: Debug,
          B: FromStr,
          B::Err: Debug,
{
    do_parse!(
        input,
        label: parse_token >>
            take_while!(is_space) >>
            set: alt!(
                do_parse!(
                    tag!("*") >> (None)
                )
                    |
                do_parse!(
                    the_set: parse_heap >> (Some(the_set))
                )
            ) >>
            (EquivalenceClass {
                label,
                set,
            })
    )
}

fn parse_heap<A>(input: &[u8]) -> IResult<&[u8], HashSet<A>>
    where A: Eq + FromStr + Hash,
          A::Err: Debug,
{
    match parse_vec(input, parse_token, "[", "]", ",") {
        IResult::Done(rest, parsed) => IResult::Done(rest, HashSet::from_iter(parsed)),
        IResult::Incomplete(needed) => IResult::Incomplete(needed),
        IResult::Error(error) => IResult::Error(error),
    }
}

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

    #[test]
    fn test_equivalence_relation_project() {
        let rel: EquivalenceRelation<u8, u8> = String::from("0 [0, 1]\n1 [2, 4]\n2 *").parse().unwrap();
        assert_eq!(0, rel.project(&1));
        assert_eq!(1, rel.project(&2));
        assert_eq!(2, rel.project(&3));
        assert_eq!(1, rel.project(&4));
    }

    #[test]
    fn test_equivalence_class_from_str_legal_input() {
        let legal_inputs = vec![
            ("0 [0, 1]", EquivalenceClass::from((0, Some(vec![0, 1])))),
            ("0  [1, 0]", EquivalenceClass::from((0, Some(vec![0, 1])))),
            ("0  [0, 1, 1]", EquivalenceClass::from((0, Some(vec![0, 1])))),
            ("0 *", EquivalenceClass::from((0, None))),
        ];

        for (legal_input, correct_class) in legal_inputs {
            assert_eq!(correct_class,
                       EquivalenceClass::from_str(legal_input).unwrap());
        }
    }

    #[test]
    fn test_equivalence_class_from_str_illegal_input() {
        let illegal_inputs = vec![
            " 0 [0, 1]",
            "[0, 1]",
            "*",
        ];

        for illegal_input in illegal_inputs {
            match EquivalenceClass::<u8, u8>::from_str(illegal_input) {
                Ok(parsed) =>
                    panic!("Was able to parse the illegal input \'{}\' as \'{:?}\'",
                           illegal_input, parsed),
                Err(_) => (),
            }
        }
    }

    #[test]
    fn test_equivalence_relation_from_str_legal_input() {
        let legal_inputs = vec![
            ("0 [0, 1]\n1 [2, 3]\n2 *", EquivalenceRelation::from(vec![
                EquivalenceClass::from((0, Some(vec![0, 1]))),
                EquivalenceClass::from((1, Some(vec![2, 3]))),
                EquivalenceClass::from((2, None)),
            ])),
        ];

        for (legal_input, correct_relation) in legal_inputs {
            assert_eq!(correct_relation,
                       EquivalenceRelation::from_str(legal_input).unwrap());
        }
    }

    #[test]
    fn test_equivalence_relation_from_str_illegal_input() {
        let illegal_inputs = vec![
            " 0 [0, 1]\n1 [2, 3]\n2 *",
            "0 [0, 1]1 [2, 3]2 *",
            "0 [0, 1]\n1 [1, 2]\n2 *",
        ];

        for illegal_input in illegal_inputs {
            match EquivalenceRelation::<u8, u8>::from_str(illegal_input) {
                Ok(parsed) =>
                    panic!("Was able to parse the illegal input \'{}\' as \'{:?}\'",
                           illegal_input, parsed),
                Err(_) => (),
            }
        }
    }
}
<commit_msg>src/approximation/equivalence_classes.rs: fixed from_str() for EquivalenceRelation<commit_after>use nom::{IResult, is_space};
use std::collections::{HashMap, HashSet};
use std::collections::hash_map::Entry;
use std::fmt::Debug;
use std::hash::Hash;
use std::iter::FromIterator;
use std::str::FromStr;

use util::parsing::*;

/// Structure containing the elements of type `A` in a equivalence class of type `B`
#[derive(Debug, Eq, PartialEq)]
pub struct EquivalenceClass<A, B>
    where A: Eq + Hash,
{
    label: B,
    set: Option<HashSet<A>>,
}

impl<A, B> From<(B, Option<Vec<A>>)> for EquivalenceClass<A, B>
    where A: Eq + Hash,
{
    fn from((label, set): (B, Option<Vec<A>>)) -> EquivalenceClass<A, B> {
        EquivalenceClass {
            label,
            set: set.and_then(|elements| Some(HashSet::from_iter(elements))),
        }
    }
}

/// A struct containing a remapping of elements of type `A` into their respective equivalence classes of type `B`.
#[derive(Clone, Debug)]
pub struct EquivalenceRelation<A, B>
    where A: Eq + Hash,
{
    map: HashMap<A, B>,
    default: B,
}

impl<A, B> PartialEq for EquivalenceRelation<A, B>
    where A: Eq + Hash,
          B: PartialEq,
{
    fn eq(&self, other: &Self) -> bool {
        self.default == other.default && self.map == other.map
    }
}

impl<A, B> EquivalenceRelation<A, B>
    where A: Eq + Hash,
          B: Clone + Eq + Hash,
{
    pub fn new(map: HashMap<B, HashSet<A>>, default: B) -> Self {
        match Self::new_safe(map, default) {
            Ok(relation) => relation,
            Err(error) => panic!(error),
        }
    }

    pub fn new_safe(map: HashMap<B, HashSet<A>>, default: B) -> Result<Self, String> {
        let mut relation_map = HashMap::new();
        for (class_name, members) in map {
            if class_name == default {
                return Err(String::from(
                    "There can only be one default class in the equivalence relation!"
                ));
            }

            for value in members {
                match relation_map.entry(value) {
                    Entry::Vacant(v) => {
                        v.insert(class_name.clone());
                    },
                    Entry::Occupied(_) => {
                        return Err(String::from(
                            "All classes of the equivalence relation must be disjoint!"
                        ));
                    },
                }
            }
        }

        Ok(EquivalenceRelation { map: relation_map, default })
    }

    // returns the equivalence class of a given value
    pub fn project(&self, key: &A) -> B {
        match self.map.get(key) {
            None => self.default.clone(),
            Some(e) => e.clone(),
        }
    }
}

impl<A, B> From<Vec<EquivalenceClass<A, B>>> for EquivalenceRelation<A, B>
    where A: Eq + Hash,
          B: Clone + Eq + Hash,
{
    fn from(classes: Vec<EquivalenceClass<A, B>>) -> EquivalenceRelation<A, B> {
        let mut map = HashMap::new();
        let mut default = None;

        for EquivalenceClass { label, set } in classes {
            if let Some(elements) = set {
                for element in elements {
                    map.insert(element, label.clone());
                }
            } else {
                default = Some(label);
            }
        }

        EquivalenceRelation { map, default: default.unwrap() }
    }
}

impl<A, B> FromStr for EquivalenceRelation<A, B>
    where A: Clone + Eq + Hash + FromStr,
          A::Err: Debug,
          B: Clone + Eq + Hash + FromStr,
          B::Err: Debug,
{
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let mut map = HashMap::new();
        let mut default = None;

        for l in s.lines() {
            if !l.is_empty() {
                match l.trim().parse()? {
                    EquivalenceClass { label, set: Some(elements) } => {
                        map.insert(label, elements);
                    },
                    EquivalenceClass { label, set: None } => {
                        default = Some(label);
                    },
                }
            }
        }

        if let Some(label) = default {
            if let Ok(relation) = EquivalenceRelation::new_safe(map, label) {
                return Ok(relation);
            }
        }

        Err(format!("Could not parse {}", s))
    }
}

impl <A, B> FromStr for EquivalenceClass<A, B>
    where A: Eq + FromStr + Hash,
          A::Err: Debug,
          B: FromStr,
          B::Err: Debug,
{
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match parse_set(s.as_bytes()) {
            IResult::Done(_, result) => Ok(result),
            _                        => Err(format!("Could not parse {}", s))
        }
    }
}

fn parse_set<A, B>(input: &[u8]) -> IResult<&[u8], EquivalenceClass<A, B>>
    where A: Eq + FromStr + Hash,
          A::Err: Debug,
          B: FromStr,
          B::Err: Debug,
{
    do_parse!(
        input,
        label: parse_token >>
            take_while!(is_space) >>
            set: alt!(
                do_parse!(
                    tag!("*") >> (None)
                )
                    |
                do_parse!(
                    the_set: parse_heap >> (Some(the_set))
                )
            ) >>
            (EquivalenceClass {
                label,
                set,
            })
    )
}

fn parse_heap<A>(input: &[u8]) -> IResult<&[u8], HashSet<A>>
    where A: Eq + FromStr + Hash,
          A::Err: Debug,
{
    match parse_vec(input, parse_token, "[", "]", ",") {
        IResult::Done(rest, parsed) => IResult::Done(rest, HashSet::from_iter(parsed)),
        IResult::Incomplete(needed) => IResult::Incomplete(needed),
        IResult::Error(error) => IResult::Error(error),
    }
}

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

    #[test]
    fn test_equivalence_relation_project() {
        let rel: EquivalenceRelation<u8, u8> = String::from("0 [0, 1]\n1 [2, 4]\n2 *").parse().unwrap();
        assert_eq!(0, rel.project(&1));
        assert_eq!(1, rel.project(&2));
        assert_eq!(2, rel.project(&3));
        assert_eq!(1, rel.project(&4));
    }

    #[test]
    fn test_equivalence_class_from_str_legal_input() {
        let legal_inputs = vec![
            ("0 [0, 1]", EquivalenceClass::from((0, Some(vec![0, 1])))),
            ("0  [1, 0]", EquivalenceClass::from((0, Some(vec![0, 1])))),
            ("0  [0, 1, 1]", EquivalenceClass::from((0, Some(vec![0, 1])))),
            ("0 *", EquivalenceClass::from((0, None))),
        ];

        for (legal_input, correct_class) in legal_inputs {
            assert_eq!(correct_class,
                       EquivalenceClass::from_str(legal_input).unwrap());
        }
    }

    #[test]
    fn test_equivalence_class_from_str_illegal_input() {
        let illegal_inputs = vec![
            " 0 [0, 1]",
            "[0, 1]",
            "*",
        ];

        for illegal_input in illegal_inputs {
            match EquivalenceClass::<u8, u8>::from_str(illegal_input) {
                Ok(parsed) =>
                    panic!("Was able to parse the illegal input \'{}\' as \'{:?}\'",
                           illegal_input, parsed),
                Err(_) => (),
            }
        }
    }

    #[test]
    fn test_equivalence_relation_from_str_legal_input() {
        let legal_inputs = vec![
            ("0 [0, 1]\n1 [2, 3]\n2 *", EquivalenceRelation::from(vec![
                EquivalenceClass::from((0, Some(vec![0, 1]))),
                EquivalenceClass::from((1, Some(vec![2, 3]))),
                EquivalenceClass::from((2, None)),
            ])),
        ];

        for (legal_input, correct_relation) in legal_inputs {
            assert_eq!(correct_relation,
                       EquivalenceRelation::from_str(legal_input).unwrap());
        }
    }

    #[test]
    fn test_equivalence_relation_from_str_illegal_input() {
        let illegal_inputs = vec![
            " 0 [0, 1]\n1 [2, 3]\n2 *",
            "0 [0, 1]1 [2, 3]2 *",
            "0 [0, 1]\n1 [1, 2]\n2 *",
        ];

        for illegal_input in illegal_inputs {
            match EquivalenceRelation::<u8, u8>::from_str(illegal_input) {
                Ok(parsed) =>
                    panic!("Was able to parse the illegal input \'{}\' as \'{:?}\'",
                           illegal_input, parsed),
                Err(_) => (),
            }
        }
    }
}
<|endoftext|>