<commit_before>use chrono::{DateTime, FixedOffset};
use std::fmt::{Display, Formatter, Result as FmtResult};
use super::deserialize_sync_user;
use model::*;

#[cfg(feature = "model")]
use std::borrow::Cow;
#[cfg(all(feature = "cache", feature = "model"))]
use CACHE;
#[cfg(all(feature = "cache", feature = "model"))]
use internal::prelude::*;
#[cfg(all(feature = "cache", feature = "model"))]
use http;
#[cfg(all(feature = "builder", feature = "cache", feature = "model"))]
use builder::EditMember;
#[cfg(all(feature = "cache", feature = "model", feature = "utils"))]
use utils::Colour;

pub trait BanOptions {
    fn dmd(&self) -> u8 { 0 }
    fn reason(&self) -> &str { "" }
}

impl BanOptions for u8 {
    fn dmd(&self) -> u8 { *self }
}

impl BanOptions for str {
    fn reason(&self) -> &str { self }
}

impl BanOptions for String {
    fn reason(&self) -> &str { self }
}

impl<'a> BanOptions for (u8, &'a str) {
    fn dmd(&self) -> u8 { self.0 }

    fn reason(&self) -> &str { self.1 }
}

impl BanOptions for (u8, String) {
    fn dmd(&self) -> u8 { self.0 }

    fn reason(&self) -> &str { &self.1 }
}

/// Information about a member of a guild.
#[derive(Clone, Debug, Deserialize)]
pub struct Member {
    /// Indicator of whether the member can hear in voice channels.
    pub deaf: bool,
    /// The unique Id of the guild that the member is a part of.
    pub guild_id: GuildId,
    /// Timestamp representing the date when the member joined.
    pub joined_at: Option<DateTime<FixedOffset>>,
    /// Indicator of whether the member can speak in voice channels.
    pub mute: bool,
    /// The member's nickname, if present.
    ///
    /// Can't be longer than 32 characters.
    pub nick: Option<String>,
    /// Vector of Ids of [`Role`]s given to the member.
    pub roles: Vec<RoleId>,
    /// Attached User struct.
    #[serde(deserialize_with = "deserialize_sync_user")]
    pub user: Arc<RwLock<User>>,
}

#[cfg(feature = "model")]
impl Member {
    /// Adds a [`Role`] to the member, editing its roles in-place if the request
    /// was successful.
    ///
    /// **Note**: Requires the [Manage Roles] permission.
    ///
    /// [`Role`]: struct.Role.html
    /// [Manage Roles]: permissions/constant.MANAGE_ROLES.html
    #[cfg(feature = "cache")]
    pub fn add_role<R: Into<RoleId>>(&mut self, role_id: R) -> Result<()> {
        let role_id = role_id.into();

        if self.roles.contains(&role_id) {
            return Ok(());
        }

        match http::add_member_role(self.guild_id.0, self.user.read().unwrap().id.0, role_id.0) {
            Ok(()) => {
                self.roles.push(role_id);

                Ok(())
            },
            Err(why) => Err(why),
        }
    }

    /// Adds one or multiple [`Role`]s to the member, editing
    /// its roles in-place if the request was successful.
    ///
    /// **Note**: Requires the [Manage Roles] permission.
    ///
    /// [`Role`]: struct.Role.html
    /// [Manage Roles]: permissions/constant.MANAGE_ROLES.html
    #[cfg(feature = "cache")]
    pub fn add_roles(&mut self, role_ids: &[RoleId]) -> Result<()> {
        self.roles.extend_from_slice(role_ids);

        let map = EditMember::default().roles(&self.roles).0;

        match http::edit_member(self.guild_id.0, self.user.read().unwrap().id.0, &map) {
            Ok(()) => Ok(()),
            Err(why) => {
                self.roles.retain(|r| !role_ids.contains(r));

                Err(why)
            },
        }
    }

    /// Ban the member from its guild, deleting the last X number of
    /// days' worth of messages.
    ///
    /// **Note**: Requires the [Ban Members] permission.
    ///
    /// # Errors
    ///
    /// Returns a [`ModelError::GuildNotFound`] if the guild could not be
    /// found.
    ///
    /// [`ModelError::GuildNotFound`]: enum.ModelError.html#variant.GuildNotFound
    ///
    /// [Ban Members]: permissions/constant.BAN_MEMBERS.html
    #[cfg(feature = "cache")]
    pub fn ban<BO: BanOptions>(&self, ban_options: BO) -> Result<()> {
        let dmd = ban_options.dmd();
        if dmd > 7 {
            return Err(Error::Model(ModelError::DeleteMessageDaysAmount(dmd)));
        }

        let reason = ban_options.reason();

        if reason.len() > 512 {
            return Err(Error::ExceededLimit(reason.to_string(), 512));
        }

        http::ban_user(
            self.guild_id.0,
            self.user.read().unwrap().id.0,
            dmd,
            &*reason,
        )
    }

    /// Determines the member's colour.
    #[cfg(all(feature = "cache", feature = "utils"))]
    pub fn colour(&self) -> Option<Colour> {
        let cache = CACHE.read().unwrap();
        let guild = match cache.guilds.get(&self.guild_id) {
            Some(guild) => guild.read().unwrap(),
            None => return None,
        };

        let mut roles = self.roles
            .iter()
            .filter_map(|role_id| guild.roles.get(role_id))
            .collect::<Vec<&Role>>();
        roles.sort_by(|a, b| b.cmp(a));

        let default = Colour::default();

        roles
            .iter()
            .find(|r| r.colour.0 != default.0)
            .map(|r| r.colour)
    }

    /// Returns the "default channel" of the guild for the member.
    /// (This returns the first channel that can be read by the member, if there isn't
    /// one returns `None`)
    pub fn default_channel(&self) -> Option<Arc<RwLock<GuildChannel>>> {
        let guild = match self.guild_id.find() {
            Some(guild) => guild,
            None => return None,
        };

        let reader = guild.read().unwrap();

        for (cid, channel) in &reader.channels {
            if reader.permissions_for(*cid, self.user.read().unwrap().id).read_messages() {
                return Some(channel.clone());
            }
        }

        None
    }

    /// Calculates the member's display name.
    ///
    /// The nickname takes priority over the member's username if it exists.
    #[inline]
    pub fn display_name(&self) -> Cow<String> {
        self.nick
            .as_ref()
            .map(Cow::Borrowed)
            .unwrap_or_else(|| Cow::Owned(self.user.read().unwrap().name.clone()))
    }

    /// Returns the DiscordTag of a Member, taking possible nickname into account.
    #[inline]
    pub fn distinct(&self) -> String {
        format!(
            "{}#{}",
            self.display_name(),
            self.user.read().unwrap().discriminator
        )
    }

    /// Edits the member with the given data. See [`Guild::edit_member`] for
    /// more information.
    ///
    /// See [`EditMember`] for the permission(s) required for separate builder
    /// methods, as well as usage of this.
    ///
    /// [`Guild::edit_member`]: ../model/struct.Guild.html#method.edit_member
    /// [`EditMember`]: ../builder/struct.EditMember.html
    #[cfg(feature = "cache")]
    pub fn edit<F: FnOnce(EditMember) -> EditMember>(&self, f: F) -> Result<()> {
        let map = f(EditMember::default()).0;

        http::edit_member(self.guild_id.0, self.user.read().unwrap().id.0, &map)
    }

    /// Kick the member from the guild.
    ///
    /// **Note**: Requires the [Kick Members] permission.
    ///
    /// # Examples
    ///
    /// Kick a member from its guild:
    ///
    /// ```rust,ignore
    /// // assuming a `member` has already been bound
    /// match member.kick() {
    ///     Ok(()) => println!("Successfully kicked member"),
    ///     Err(Error::Model(ModelError::GuildNotFound)) => {
    ///         println!("Couldn't determine guild of member");
    ///     },
    ///     Err(Error::Model(ModelError::InvalidPermissions(missing_perms))) => {
    ///         println!("Didn't have permissions; missing: {:?}", missing_perms);
    ///     },
    ///     _ => {},
    /// }
    /// ```
    ///
    /// # Errors
    ///
    /// Returns a [`ModelError::GuildNotFound`] if the Id of the member's guild
    /// could not be determined.
    ///
    /// If the `cache` is enabled, returns a [`ModelError::InvalidPermissions`]
    /// if the current user does not have permission to perform the kick.
    ///
    /// [`ModelError::GuildNotFound`]: enum.ModelError.html#variant.GuildNotFound
    /// [`ModelError::InvalidPermissions`]: enum.ModelError.html#variant.InvalidPermissions
    /// [Kick Members]: permissions/constant.KICK_MEMBERS.html
    pub fn kick(&self) -> Result<()> {
        #[cfg(feature = "cache")]
        {
            let req = Permissions::KICK_MEMBERS;

            let has_perms = CACHE
                .read()
                .unwrap()
                .guilds
                .get(&self.guild_id)
                .map(|guild| guild.read().unwrap().has_perms(req));

            if let Some(Ok(false)) = has_perms {
                return Err(Error::Model(ModelError::InvalidPermissions(req)));
            }
        }

        self.guild_id.kick(self.user.read().unwrap().id)
    }

    /// Returns the permissions for the member.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// // assuming there's a `member` variable gotten from anything.
    /// println!("The permission bits for the member are: {}",
    /// member.permissions().expect("permissions").bits);
    /// ```
    ///
    /// # Errors
    ///
    /// Returns a [`ModelError::GuildNotFound`] if the guild the member's in could not be
    /// found in the cache.
    ///
    /// And/or returns [`ModelError::ItemMissing`] if the "default channel" of the guild is not
    /// found.
    ///
    /// [`ModelError::GuildNotFound`]: enum.ModelError.html#variant.GuildNotFound
    /// [`ModelError::ItemMissing`]: enum.ModelError.html#variant.ItemMissing
    #[cfg(feature = "cache")]
    pub fn permissions(&self) -> Result<Permissions> {
        let guild = match self.guild_id.find() {
            Some(guild) => guild,
            None => return Err(From::from(ModelError::GuildNotFound)),
        };

        let guild = guild.read().unwrap();

        let default_channel = match guild.default_channel(self.user.read().unwrap().id) {
            Some(dc) => dc,
            None => return Err(From::from(ModelError::ItemMissing)),
        };

        let default_channel_reader = default_channel.read().unwrap();

        Ok(
            guild
                .permissions_for(default_channel_reader.id, self.user.read().unwrap().id),
        )
    }

    /// Removes a [`Role`] from the member, editing its roles in-place if the
    /// request was successful.
    ///
    /// **Note**: Requires the [Manage Roles] permission.
    ///
    /// [`Role`]: struct.Role.html
    /// [Manage Roles]: permissions/constant.MANAGE_ROLES.html
    #[cfg(feature = "cache")]
    pub fn remove_role<R: Into<RoleId>>(&mut self, role_id: R) -> Result<()> {
        let role_id = role_id.into();

        if !self.roles.contains(&role_id) {
            return Ok(());
        }

        match http::remove_member_role(self.guild_id.0, self.user.read().unwrap().id.0, role_id.0) {
            Ok(()) => {
                self.roles.retain(|r| r.0 != role_id.0);

                Ok(())
            },
            Err(why) => Err(why),
        }
    }

    /// Removes one or multiple [`Role`]s from the member.
    ///
    /// **Note**: Requires the [Manage Roles] permission.
    ///
    /// [`Role`]: struct.Role.html
    /// [Manage Roles]: permissions/constant.MANAGE_ROLES.html
    #[cfg(feature = "cache")]
    pub fn remove_roles(&mut self, role_ids: &[RoleId]) -> Result<()> {
        self.roles.retain(|r| !role_ids.contains(r));

        let map = EditMember::default().roles(&self.roles).0;

        match http::edit_member(self.guild_id.0, self.user.read().unwrap().id.0, &map) {
            Ok(()) => Ok(()),
            Err(why) => {
                self.roles.extend_from_slice(role_ids);

                Err(why)
            },
        }
    }

    /// Retrieves the full role data for the user's roles.
    ///
    /// This is shorthand for manually searching through the CACHE.
    ///
    /// If role data can not be found for the member, then `None` is returned.
    #[cfg(feature = "cache")]
    pub fn roles(&self) -> Option<Vec<Role>> {
        CACHE
            .read()
            .unwrap()
            .guilds
            .values()
            .find(|guild| {
                guild.read().unwrap().members.values().any(|m| {
                    m.user.read().unwrap().id == self.user.read().unwrap().id &&
                    m.joined_at == self.joined_at
                })
            })
            .map(|guild| {
                guild
                    .read()
                    .unwrap()
                    .roles
                    .values()
                    .filter(|role| self.roles.contains(&role.id))
                    .cloned()
                    .collect()
            })
    }

    /// Unbans the [`User`] from the guild.
    ///
    /// **Note**: Requires the [Ban Members] permission.
    ///
    /// # Errors
    ///
    /// If the `cache` is enabled, returns a [`ModelError::InvalidPermissions`]
    /// if the current user does not have permission to perform bans.
    ///
    /// [`ModelError::InvalidPermissions`]: enum.ModelError.html#variant.InvalidPermissions
    /// [`User`]: struct.User.html
    /// [Ban Members]: permissions/constant.BAN_MEMBERS.html
    #[cfg(feature = "cache")]
    pub fn unban(&self) -> Result<()> {
        http::remove_ban(self.guild_id.0, self.user.read().unwrap().id.0)
    }
}

impl Display for Member {
    /// Mentions the user so that they receive a notification.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// // assumes a `member` has already been bound
    /// println!("{} is a member!", member);
    /// ```
    ///
    // This is in the format of `<@USER_ID>`.
    fn fmt(&self, f: &mut Formatter) -> FmtResult {
        Display::fmt(&self.user.read().unwrap().mention(), f)
    }
}
<commit_msg>Add an impl for `&str`<commit_after>use chrono::{DateTime, FixedOffset};
use std::fmt::{Display, Formatter, Result as FmtResult};
use super::deserialize_sync_user;
use model::*;

#[cfg(feature = "model")]
use std::borrow::Cow;
#[cfg(all(feature = "cache", feature = "model"))]
use CACHE;
#[cfg(all(feature = "cache", feature = "model"))]
use internal::prelude::*;
#[cfg(all(feature = "cache", feature = "model"))]
use http;
#[cfg(all(feature = "builder", feature = "cache", feature = "model"))]
use builder::EditMember;
#[cfg(all(feature = "cache", feature = "model", feature = "utils"))]
use utils::Colour;

pub trait BanOptions {
    fn dmd(&self) -> u8 { 0 }
    fn reason(&self) -> &str { "" }
}

impl BanOptions for u8 {
    fn dmd(&self) -> u8 { *self }
}

impl BanOptions for str {
    fn reason(&self) -> &str { self }
}

impl<'a> BanOptions for &'a str {
    fn reason(&self) -> &str { self }
}

impl BanOptions for String {
    fn reason(&self) -> &str { self }
}

impl<'a> BanOptions for (u8, &'a str) {
    fn dmd(&self) -> u8 { self.0 }

    fn reason(&self) -> &str { self.1 }
}

impl BanOptions for (u8, String) {
    fn dmd(&self) -> u8 { self.0 }

    fn reason(&self) -> &str { &self.1 }
}

/// Information about a member of a guild.
#[derive(Clone, Debug, Deserialize)]
pub struct Member {
    /// Indicator of whether the member can hear in voice channels.
    pub deaf: bool,
    /// The unique Id of the guild that the member is a part of.
    pub guild_id: GuildId,
    /// Timestamp representing the date when the member joined.
    pub joined_at: Option<DateTime<FixedOffset>>,
    /// Indicator of whether the member can speak in voice channels.
    pub mute: bool,
    /// The member's nickname, if present.
    ///
    /// Can't be longer than 32 characters.
    pub nick: Option<String>,
    /// Vector of Ids of [`Role`]s given to the member.
    pub roles: Vec<RoleId>,
    /// Attached User struct.
    #[serde(deserialize_with = "deserialize_sync_user")]
    pub user: Arc<RwLock<User>>,
}

#[cfg(feature = "model")]
impl Member {
    /// Adds a [`Role`] to the member, editing its roles in-place if the request
    /// was successful.
    ///
    /// **Note**: Requires the [Manage Roles] permission.
    ///
    /// [`Role`]: struct.Role.html
    /// [Manage Roles]: permissions/constant.MANAGE_ROLES.html
    #[cfg(feature = "cache")]
    pub fn add_role<R: Into<RoleId>>(&mut self, role_id: R) -> Result<()> {
        let role_id = role_id.into();

        if self.roles.contains(&role_id) {
            return Ok(());
        }

        match http::add_member_role(self.guild_id.0, self.user.read().unwrap().id.0, role_id.0) {
            Ok(()) => {
                self.roles.push(role_id);

                Ok(())
            },
            Err(why) => Err(why),
        }
    }

    /// Adds one or multiple [`Role`]s to the member, editing
    /// its roles in-place if the request was successful.
    ///
    /// **Note**: Requires the [Manage Roles] permission.
    ///
    /// [`Role`]: struct.Role.html
    /// [Manage Roles]: permissions/constant.MANAGE_ROLES.html
    #[cfg(feature = "cache")]
    pub fn add_roles(&mut self, role_ids: &[RoleId]) -> Result<()> {
        self.roles.extend_from_slice(role_ids);

        let map = EditMember::default().roles(&self.roles).0;

        match http::edit_member(self.guild_id.0, self.user.read().unwrap().id.0, &map) {
            Ok(()) => Ok(()),
            Err(why) => {
                self.roles.retain(|r| !role_ids.contains(r));

                Err(why)
            },
        }
    }

    /// Ban the member from its guild, deleting the last X number of
    /// days' worth of messages.
    ///
    /// **Note**: Requires the [Ban Members] permission.
    ///
    /// # Errors
    ///
    /// Returns a [`ModelError::GuildNotFound`] if the guild could not be
    /// found.
    ///
    /// [`ModelError::GuildNotFound`]: enum.ModelError.html#variant.GuildNotFound
    ///
    /// [Ban Members]: permissions/constant.BAN_MEMBERS.html
    #[cfg(feature = "cache")]
    pub fn ban<BO: BanOptions>(&self, ban_options: BO) -> Result<()> {
        let dmd = ban_options.dmd();
        if dmd > 7 {
            return Err(Error::Model(ModelError::DeleteMessageDaysAmount(dmd)));
        }

        let reason = ban_options.reason();

        if reason.len() > 512 {
            return Err(Error::ExceededLimit(reason.to_string(), 512));
        }

        http::ban_user(
            self.guild_id.0,
            self.user.read().unwrap().id.0,
            dmd,
            &*reason,
        )
    }

    /// Determines the member's colour.
    #[cfg(all(feature = "cache", feature = "utils"))]
    pub fn colour(&self) -> Option<Colour> {
        let cache = CACHE.read().unwrap();
        let guild = match cache.guilds.get(&self.guild_id) {
            Some(guild) => guild.read().unwrap(),
            None => return None,
        };

        let mut roles = self.roles
            .iter()
            .filter_map(|role_id| guild.roles.get(role_id))
            .collect::<Vec<&Role>>();
        roles.sort_by(|a, b| b.cmp(a));

        let default = Colour::default();

        roles
            .iter()
            .find(|r| r.colour.0 != default.0)
            .map(|r| r.colour)
    }

    /// Returns the "default channel" of the guild for the member.
    /// (This returns the first channel that can be read by the member, if there isn't
    /// one returns `None`)
    pub fn default_channel(&self) -> Option<Arc<RwLock<GuildChannel>>> {
        let guild = match self.guild_id.find() {
            Some(guild) => guild,
            None => return None,
        };

        let reader = guild.read().unwrap();

        for (cid, channel) in &reader.channels {
            if reader.permissions_for(*cid, self.user.read().unwrap().id).read_messages() {
                return Some(channel.clone());
            }
        }

        None
    }

    /// Calculates the member's display name.
    ///
    /// The nickname takes priority over the member's username if it exists.
    #[inline]
    pub fn display_name(&self) -> Cow<String> {
        self.nick
            .as_ref()
            .map(Cow::Borrowed)
            .unwrap_or_else(|| Cow::Owned(self.user.read().unwrap().name.clone()))
    }

    /// Returns the DiscordTag of a Member, taking possible nickname into account.
    #[inline]
    pub fn distinct(&self) -> String {
        format!(
            "{}#{}",
            self.display_name(),
            self.user.read().unwrap().discriminator
        )
    }

    /// Edits the member with the given data. See [`Guild::edit_member`] for
    /// more information.
    ///
    /// See [`EditMember`] for the permission(s) required for separate builder
    /// methods, as well as usage of this.
    ///
    /// [`Guild::edit_member`]: ../model/struct.Guild.html#method.edit_member
    /// [`EditMember`]: ../builder/struct.EditMember.html
    #[cfg(feature = "cache")]
    pub fn edit<F: FnOnce(EditMember) -> EditMember>(&self, f: F) -> Result<()> {
        let map = f(EditMember::default()).0;

        http::edit_member(self.guild_id.0, self.user.read().unwrap().id.0, &map)
    }

    /// Kick the member from the guild.
    ///
    /// **Note**: Requires the [Kick Members] permission.
    ///
    /// # Examples
    ///
    /// Kick a member from its guild:
    ///
    /// ```rust,ignore
    /// // assuming a `member` has already been bound
    /// match member.kick() {
    ///     Ok(()) => println!("Successfully kicked member"),
    ///     Err(Error::Model(ModelError::GuildNotFound)) => {
    ///         println!("Couldn't determine guild of member");
    ///     },
    ///     Err(Error::Model(ModelError::InvalidPermissions(missing_perms))) => {
    ///         println!("Didn't have permissions; missing: {:?}", missing_perms);
    ///     },
    ///     _ => {},
    /// }
    /// ```
    ///
    /// # Errors
    ///
    /// Returns a [`ModelError::GuildNotFound`] if the Id of the member's guild
    /// could not be determined.
    ///
    /// If the `cache` is enabled, returns a [`ModelError::InvalidPermissions`]
    /// if the current user does not have permission to perform the kick.
    ///
    /// [`ModelError::GuildNotFound`]: enum.ModelError.html#variant.GuildNotFound
    /// [`ModelError::InvalidPermissions`]: enum.ModelError.html#variant.InvalidPermissions
    /// [Kick Members]: permissions/constant.KICK_MEMBERS.html
    pub fn kick(&self) -> Result<()> {
        #[cfg(feature = "cache")]
        {
            let req = Permissions::KICK_MEMBERS;

            let has_perms = CACHE
                .read()
                .unwrap()
                .guilds
                .get(&self.guild_id)
                .map(|guild| guild.read().unwrap().has_perms(req));

            if let Some(Ok(false)) = has_perms {
                return Err(Error::Model(ModelError::InvalidPermissions(req)));
            }
        }

        self.guild_id.kick(self.user.read().unwrap().id)
    }

    /// Returns the permissions for the member.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// // assuming there's a `member` variable gotten from anything.
    /// println!("The permission bits for the member are: {}",
    /// member.permissions().expect("permissions").bits);
    /// ```
    ///
    /// # Errors
    ///
    /// Returns a [`ModelError::GuildNotFound`] if the guild the member's in could not be
    /// found in the cache.
    ///
    /// And/or returns [`ModelError::ItemMissing`] if the "default channel" of the guild is not
    /// found.
    ///
    /// [`ModelError::GuildNotFound`]: enum.ModelError.html#variant.GuildNotFound
    /// [`ModelError::ItemMissing`]: enum.ModelError.html#variant.ItemMissing
    #[cfg(feature = "cache")]
    pub fn permissions(&self) -> Result<Permissions> {
        let guild = match self.guild_id.find() {
            Some(guild) => guild,
            None => return Err(From::from(ModelError::GuildNotFound)),
        };

        let guild = guild.read().unwrap();

        let default_channel = match guild.default_channel(self.user.read().unwrap().id) {
            Some(dc) => dc,
            None => return Err(From::from(ModelError::ItemMissing)),
        };

        let default_channel_reader = default_channel.read().unwrap();

        Ok(
            guild
                .permissions_for(default_channel_reader.id, self.user.read().unwrap().id),
        )
    }

    /// Removes a [`Role`] from the member, editing its roles in-place if the
    /// request was successful.
    ///
    /// **Note**: Requires the [Manage Roles] permission.
    ///
    /// [`Role`]: struct.Role.html
    /// [Manage Roles]: permissions/constant.MANAGE_ROLES.html
    #[cfg(feature = "cache")]
    pub fn remove_role<R: Into<RoleId>>(&mut self, role_id: R) -> Result<()> {
        let role_id = role_id.into();

        if !self.roles.contains(&role_id) {
            return Ok(());
        }

        match http::remove_member_role(self.guild_id.0, self.user.read().unwrap().id.0, role_id.0) {
            Ok(()) => {
                self.roles.retain(|r| r.0 != role_id.0);

                Ok(())
            },
            Err(why) => Err(why),
        }
    }

    /// Removes one or multiple [`Role`]s from the member.
    ///
    /// **Note**: Requires the [Manage Roles] permission.
    ///
    /// [`Role`]: struct.Role.html
    /// [Manage Roles]: permissions/constant.MANAGE_ROLES.html
    #[cfg(feature = "cache")]
    pub fn remove_roles(&mut self, role_ids: &[RoleId]) -> Result<()> {
        self.roles.retain(|r| !role_ids.contains(r));

        let map = EditMember::default().roles(&self.roles).0;

        match http::edit_member(self.guild_id.0, self.user.read().unwrap().id.0, &map) {
            Ok(()) => Ok(()),
            Err(why) => {
                self.roles.extend_from_slice(role_ids);

                Err(why)
            },
        }
    }

    /// Retrieves the full role data for the user's roles.
    ///
    /// This is shorthand for manually searching through the CACHE.
    ///
    /// If role data can not be found for the member, then `None` is returned.
    #[cfg(feature = "cache")]
    pub fn roles(&self) -> Option<Vec<Role>> {
        CACHE
            .read()
            .unwrap()
            .guilds
            .values()
            .find(|guild| {
                guild.read().unwrap().members.values().any(|m| {
                    m.user.read().unwrap().id == self.user.read().unwrap().id &&
                    m.joined_at == self.joined_at
                })
            })
            .map(|guild| {
                guild
                    .read()
                    .unwrap()
                    .roles
                    .values()
                    .filter(|role| self.roles.contains(&role.id))
                    .cloned()
                    .collect()
            })
    }

    /// Unbans the [`User`] from the guild.
    ///
    /// **Note**: Requires the [Ban Members] permission.
    ///
    /// # Errors
    ///
    /// If the `cache` is enabled, returns a [`ModelError::InvalidPermissions`]
    /// if the current user does not have permission to perform bans.
    ///
    /// [`ModelError::InvalidPermissions`]: enum.ModelError.html#variant.InvalidPermissions
    /// [`User`]: struct.User.html
    /// [Ban Members]: permissions/constant.BAN_MEMBERS.html
    #[cfg(feature = "cache")]
    pub fn unban(&self) -> Result<()> {
        http::remove_ban(self.guild_id.0, self.user.read().unwrap().id.0)
    }
}

impl Display for Member {
    /// Mentions the user so that they receive a notification.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// // assumes a `member` has already been bound
    /// println!("{} is a member!", member);
    /// ```
    ///
    // This is in the format of `<@USER_ID>`.
    fn fmt(&self, f: &mut Formatter) -> FmtResult {
        Display::fmt(&self.user.read().unwrap().mention(), f)
    }
}
<|endoftext|>


<commit_before>use chrono::DateTime;
use chrono::offset::Utc;
use diesel;
use diesel::prelude::*;
use diesel::pg::PgConnection;
use diesel::pg::upsert::*;
use futures::{Future, stream, Stream};
use futures_cpupool::CpuPool;
use hyper::{Body, Client, Uri};
use hyper::client::HttpConnector;
use hyper_tls::HttpsConnector;
use InsertError;
use kuchiki;
use kuchiki::traits::*;
use r2d2_diesel::ConnectionManager;
use r2d2::Pool;
use rusoto_core::{default_tls_client, Region};
use rusoto_credential::DefaultCredentialsProvider;
use rusoto_s3::{PutObjectRequest, S3Client, S3};
use schema::ads;
use server::AdPost;

const ENDPOINT: &'static str = "https://pp-facebook-ads.s3.amazonaws.com/";

fn get_title(document: &kuchiki::NodeRef) -> Result<String, InsertError> {
    document
        .select("h5 a, h6 a, strong")
        .map_err(InsertError::HTML)?
        .nth(0)
        .and_then(|a| Some(a.text_contents()))
        .ok_or(InsertError::HTML(()))
}

fn get_image(document: &kuchiki::NodeRef) -> Result<String, InsertError> {
    document
        .select("img")
        .map_err(InsertError::HTML)?
        .nth(0)
        .and_then(|a| {
            a.attributes.borrow().get("src").and_then(
                |src| Some(src.to_string()),
            )
        })
        .ok_or(InsertError::HTML(()))
}

fn get_message(document: &kuchiki::NodeRef) -> Result<String, InsertError> {
    let selectors = vec![".userContent p", "span"];
    let iters = selectors.iter().map(|s| document.select(s)).flat_map(|a| a);

    iters
        .map(|i| {
            i.fold(String::new(), |m, a| m + &a.as_node().to_string())
        })
        .filter(|i| !i.is_empty())
        .nth(0)
        .ok_or(InsertError::HTML(()))
}

fn get_images(document: &kuchiki::NodeRef) -> Result<Vec<String>, InsertError> {
    Ok(
        document
            .select("img")
            .map_err(InsertError::HTML)?
            .skip(1)
            .map(|a| {
                a.attributes.borrow().get("src").and_then(
                    |s| Some(s.to_string()),
                )
            })
            .filter(|s| s.is_some())
            .map(|s| s.unwrap())
            .collect::<Vec<String>>(),
    )
}

#[derive(AsChangeset, Debug)]
#[table_name = "ads"]
pub struct Images {
    thumbnail: Option<String>,
    images: Vec<String>,
    title: String,
    message: String,
    html: String,
}

impl Images {
    fn from_ad(ad: &Ad, images: Vec<Uri>) -> Result<Images, InsertError> {
        let thumb = images
            .iter()
            .filter(|i| ad.thumbnail.contains(i.path()))
            .map(|i| ENDPOINT.to_string() + i.path().trim_left_matches('/'))
            .nth(0);

        let mut rest = images.clone();
        if let Some(thumb) = thumb.clone() {
            rest.retain(|x| !thumb.contains(x.path()))
        };

        let collection = rest.iter()
            .filter(|i| ad.images.iter().any(|a| a.contains(i.path())))
            .map(|i| ENDPOINT.to_string() + i.path().trim_left_matches('/'))
            .collect::<Vec<String>>();

        let document = kuchiki::parse_html().one(ad.html.clone());
        for a in document.select("img").map_err(InsertError::HTML)? {
            if let Some(x) = a.attributes.borrow_mut().get_mut("src") {
                if let Ok(u) = x.parse::<Uri>() {
                    if let Some(i) = images.iter().find(|i| i.path() == u.path()) {
                        *x = ENDPOINT.to_string() + i.path().trim_left_matches('/');
                    } else {
                        *x = "".to_string();
                    }
                }
            };
        }

        let title = get_title(&document)?;
        let message = get_title(&document)?;
        Ok(Images {
            thumbnail: thumb,
            images: collection,
            title: title,
            html: document
                .select("div")
                .map_err(InsertError::HTML)?
                .nth(0)
                .ok_or(InsertError::HTML(()))?
                .as_node()
                .to_string(),
            message: message,
        })
    }
}

#[derive(Serialize, Queryable, Debug, Clone)]
pub struct Ad {
    pub id: String,
    pub html: String,
    pub political: i32,
    pub not_political: i32,
    pub title: String,
    pub message: String,
    pub thumbnail: String,
    pub created_at: DateTime<Utc>,
    pub updated_at: DateTime<Utc>,
    pub lang: String,
    pub images: Vec<String>,
    pub impressions: i32,
    pub political_probability: f64,
    pub targeting: Option<String>,
}

impl Ad {
    // This will asynchronously save the images to s3 we may very well end up
    // dropping images, but I can't see any way around it right now. Also we
    // should think about splitting this up, but I'm fine -- if a little
    // embarassed about it -- right now. This function swallows errors, and
    // there's a chance we'll end up with no images at the end, but I think we
    // can handle that in the extension's UI.
    pub fn grab_and_store(
        &self,
        client: Client<HttpsConnector<HttpConnector>, Body>,
        db: &Pool<ConnectionManager<PgConnection>>,
        pool: CpuPool,
    ) -> Box<Future<Item = (), Error = ()>> {
        let ad = self.clone();
        let pool_s3 = pool.clone();
        let pool_db = pool.clone();
        let db = db.clone();
        let future = stream::iter_ok(self.image_urls())
            .map(|u| u.unwrap())
            // filter ones we already have in the db and ones we can verify as
            // coming from fb, we don't want to become a malware vector :)
            // currently we redownload images we already have, but ok.
            .filter(|u| {
                info!("testing {:?}", u.host());
                match u.host() {
                    Some(h) => (h == "pp-facebook-ads.s3.amazonaws.com" || h.ends_with("fbcdn.net")),
                    None => false
                }
            })
            // grab image
            .and_then(move |img| {
                info!("getting {:?}", img.path());
                let cloned = img.clone();
                client
                    .get(img)
                    .and_then(|res| {
                        res.body().concat2().and_then(|chunk| Ok((chunk, cloned)))
                    })
                    .map_err(InsertError::Hyper)
            })
            // upload them to s3
            .and_then(move |tuple| {
                let pool = pool_s3.clone();
                // we do this in a worker thread because rusoto isn't on
                // Hyper async yet.
                pool.spawn_fn(move || {
                    if tuple.1.host().unwrap() != "pp-facebook-ads.s3.amazonaws.com" {
                        let credentials = DefaultCredentialsProvider::new()
                            .map_err(InsertError::AWS)?;
                        let tls = default_tls_client()
                            .map_err(InsertError::TLS)?;
                        let client = S3Client::new(tls, credentials, Region::UsEast1);
                        let req = PutObjectRequest {
                            bucket: "pp-facebook-ads".to_string(),
                            key: tuple.1.path().trim_left_matches('/').to_string(),
                            acl: Some("public-read".to_string()),
                            body: Some(tuple.0.to_vec()),
                            ..PutObjectRequest::default()
                        };
                        client.put_object(&req).map_err(InsertError::S3)?;
                    }
                    Ok(tuple.1)
                })
            })
            .collect()
            // save the new urls to the database. the images variable will
            // include only those that we've successfully saved to s3, so we
            // have to do a funky merge here.
            .and_then(move |images| {
                let imgs = images.clone();
                pool_db.spawn_fn(move || {
                    use schema::ads::dsl::*;
                    let update = Images::from_ad(&ad, imgs)?;
                    let connection = db.get().map_err(InsertError::Timeout)?;
                    diesel::update(ads.filter(id.eq(&ad.id)))
                        .set(&update)
                        .execute(&*connection)
                        .map_err(InsertError::DataBase)?;
                    info!("saved {:?}", ad.id);
                    Ok(())
                })
            })
            .map_err(|e| {
                warn!("{:?}", e);
                ()
            });
        Box::new(future)
    }

    pub(self) fn image_urls(&self) -> Vec<Result<Uri, InsertError>> {
        let images = [vec![self.thumbnail.clone()], self.images.clone()];
        images
            .concat()
            .iter()
            .map(|a| {
                let a: Uri = a.parse().map_err(InsertError::Uri)?;
                Ok(a)
            })
            .collect()
    }

    pub fn get_ads_by_lang(
        language: &str,
        conn: &Pool<ConnectionManager<PgConnection>>,
    ) -> Result<Vec<Ad>, InsertError> {
        use schema::ads::dsl::*;
        let connection = conn.get().map_err(InsertError::Timeout)?;
        ads.filter(lang.eq(language))
            .filter(political_probability.gt(0.75))
            .order(created_at.desc())
            .limit(100)
            .load::<Ad>(&*connection)
            .map_err(InsertError::DataBase)
    }
}

#[derive(Insertable)]
#[table_name = "ads"]
pub struct NewAd<'a> {
    id: &'a str,
    html: &'a str,
    political: i32,
    not_political: i32,

    title: String,
    message: String,
    thumbnail: String,

    lang: &'a str,
    images: Vec<String>,
    impressions: i32,

    targeting: Option<String>,
}

impl<'a> NewAd<'a> {
    pub fn new(ad: &'a AdPost, lang: &'a str) -> Result<NewAd<'a>, InsertError> {
        info!("saving {}", ad.id);
        let document = kuchiki::parse_html().one(ad.html.clone());

        let thumb = get_image(&document)?;
        let images = get_images(&document)?;
        let message = get_message(&document)?;
        let title = get_title(&document)?;

        Ok(NewAd {
            id: &ad.id,
            html: &ad.html,
            // we try unwrapping or we chose the false branch in both of these
            // cases to count impressions
            political: if ad.political.unwrap_or(false) { 1 } else { 0 },
            not_political: if !ad.political.unwrap_or(true) { 1 } else { 0 },
            title: title,
            message: message,
            thumbnail: thumb,
            lang: lang,
            images: images,
            impressions: if !ad.political.is_some() { 1 } else { 0 },
            targeting: ad.targeting.clone(),
        })
    }

    pub fn save(&self, pool: &Pool<ConnectionManager<PgConnection>>) -> Result<Ad, InsertError> {
        use schema::ads;
        use schema::ads::dsl::*;
        let connection = pool.get().map_err(InsertError::Timeout)?;

        // increment impressions if this is a background save,
        // otherwise increment political counters
        let ad: Ad = diesel::insert(&self.on_conflict(
            id,
            do_update().set((
                political.eq(political + self.political),
                not_political.eq(
                    not_political +
                        self.not_political,
                ),
                impressions.eq(
                    impressions + self.impressions,
                ),
                updated_at.eq(Utc::now()),
            )),
        )).into(ads::table)
            .get_result(&*connection)
            .map_err(InsertError::DataBase)?;

        if self.targeting.is_some() && !ad.targeting.is_some() {
            diesel::update(ads.filter(id.eq(self.id)))
                .set(targeting.eq(self.targeting.clone()))
                .execute(&*connection)
                .map_err(InsertError::DataBase)?;
        };

        Ok(ad)
    }
}

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

    #[test]
    fn ad_parsing() {
        let ad = include_str!("./html-test.txt");
        let post = AdPost {
            id: "test".to_string(),
            html: ad.to_string(),
            political: None,
            targeting: None,
        };
        let new_ad = NewAd::new(&post, "en-US").unwrap();
        assert!(new_ad.thumbnail.len() > 0);
        assert_eq!(new_ad.images.len(), 2);
        assert!(new_ad.title.len() > 0);

        assert_eq!(
            kuchiki::parse_html().one(new_ad.message).text_contents(),
            kuchiki::parse_html().one("
#Valerian is “the best experience since ‘Avatar.’” See it in 3D and RealD3D theaters this Friday. Get tickets now: ValerianTickets.com
").text_contents()
        );
    }

    #[test]
    fn image_parsing() {
        use chrono::prelude::*;
        let ad = include_str!("./html-test.txt");
        let document = kuchiki::parse_html().one(ad);

        let saved_ad = Ad {
            id: "test".to_string(),
            html: ad.to_string(),
            political: 1,
            not_political: 2,
            title: get_title(&document).unwrap(),
            message: get_message(&document).unwrap(),
            thumbnail: get_image(&document).unwrap(),
            created_at: Utc::now(),
            updated_at: Utc::now(),
            lang: "US".to_string(),
            images: get_images(&document).unwrap(),
            impressions: 1,
            targeting: None,
            political_probability: 0.0,
        };
        let urls = saved_ad
            .image_urls()
            .into_iter()
            .map(|x| x.unwrap())
            .collect();
        let images = Images::from_ad(&saved_ad, urls).unwrap();
        assert!(images.html != saved_ad.html);
        assert!(!images.html.contains("fbcdn"));
        assert!(!images.html.contains("html"));
        assert!(images.images.len() == saved_ad.images.len());
        assert!(images.thumbnail.unwrap() != saved_ad.thumbnail);
    }
}
<commit_msg>fix image parsing<commit_after>use chrono::DateTime;
use chrono::offset::Utc;
use diesel;
use diesel::prelude::*;
use diesel::pg::PgConnection;
use diesel::pg::upsert::*;
use futures::{Future, stream, Stream};
use futures_cpupool::CpuPool;
use hyper::{Body, Client, Uri};
use hyper::client::HttpConnector;
use hyper_tls::HttpsConnector;
use InsertError;
use kuchiki;
use kuchiki::traits::*;
use r2d2_diesel::ConnectionManager;
use r2d2::Pool;
use rusoto_core::{default_tls_client, Region};
use rusoto_credential::DefaultCredentialsProvider;
use rusoto_s3::{PutObjectRequest, S3Client, S3};
use schema::ads;
use server::AdPost;

const ENDPOINT: &'static str = "https://pp-facebook-ads.s3.amazonaws.com/";

fn get_title(document: &kuchiki::NodeRef) -> Result<String, InsertError> {
    document
        .select("h5 a, h6 a, strong")
        .map_err(InsertError::HTML)?
        .nth(0)
        .and_then(|a| Some(a.text_contents()))
        .ok_or(InsertError::HTML(()))
}

fn get_image(document: &kuchiki::NodeRef) -> Result<String, InsertError> {
    document
        .select("img")
        .map_err(InsertError::HTML)?
        .nth(0)
        .and_then(|a| {
            a.attributes.borrow().get("src").and_then(
                |src| Some(src.to_string()),
            )
        })
        .ok_or(InsertError::HTML(()))
}

fn get_message(document: &kuchiki::NodeRef) -> Result<String, InsertError> {
    let selectors = vec![".userContent p", "span"];
    let iters = selectors.iter().map(|s| document.select(s)).flat_map(|a| a);

    iters
        .map(|i| {
            i.fold(String::new(), |m, a| m + &a.as_node().to_string())
        })
        .filter(|i| !i.is_empty())
        .nth(0)
        .ok_or(InsertError::HTML(()))
}

fn get_images(document: &kuchiki::NodeRef) -> Result<Vec<String>, InsertError> {
    Ok(
        document
            .select("img")
            .map_err(InsertError::HTML)?
            .skip(1)
            .map(|a| {
                a.attributes.borrow().get("src").and_then(
                    |s| Some(s.to_string()),
                )
            })
            .filter(|s| s.is_some())
            .map(|s| s.unwrap())
            .collect::<Vec<String>>(),
    )
}

#[derive(AsChangeset, Debug)]
#[table_name = "ads"]
pub struct Images {
    thumbnail: Option<String>,
    images: Vec<String>,
    title: String,
    message: String,
    html: String,
}

impl Images {
    fn from_ad(ad: &Ad, images: Vec<Uri>) -> Result<Images, InsertError> {
        let thumb = images
            .iter()
            .filter(|i| ad.thumbnail.contains(i.path()))
            .map(|i| ENDPOINT.to_string() + i.path().trim_left_matches('/'))
            .nth(0);

        let mut rest = images.clone();
        if let Some(thumb) = thumb.clone() {
            rest.retain(|x| !thumb.contains(x.path()))
        };

        let collection = rest.iter()
            .filter(|i| ad.images.iter().any(|a| a.contains(i.path())))
            .map(|i| ENDPOINT.to_string() + i.path().trim_left_matches('/'))
            .collect::<Vec<String>>();

        let document = kuchiki::parse_html().one(ad.html.clone());
        for a in document.select("img").map_err(InsertError::HTML)? {
            if let Some(x) = a.attributes.borrow_mut().get_mut("src") {
                if let Ok(u) = x.parse::<Uri>() {
                    if let Some(i) = images.iter().find(|i| i.path() == u.path()) {
                        *x = ENDPOINT.to_string() + i.path().trim_left_matches('/');
                    } else {
                        *x = "".to_string();
                    }
                }
            };
        }

        let title = get_title(&document)?;
        let message = get_message(&document)?;
        Ok(Images {
            thumbnail: thumb,
            images: collection,
            title: title,
            html: document
                .select("div")
                .map_err(InsertError::HTML)?
                .nth(0)
                .ok_or(InsertError::HTML(()))?
                .as_node()
                .to_string(),
            message: message,
        })
    }
}

#[derive(Serialize, Queryable, Debug, Clone)]
pub struct Ad {
    pub id: String,
    pub html: String,
    pub political: i32,
    pub not_political: i32,
    pub title: String,
    pub message: String,
    pub thumbnail: String,
    pub created_at: DateTime<Utc>,
    pub updated_at: DateTime<Utc>,
    pub lang: String,
    pub images: Vec<String>,
    pub impressions: i32,
    pub political_probability: f64,
    pub targeting: Option<String>,
}

impl Ad {
    // This will asynchronously save the images to s3 we may very well end up
    // dropping images, but I can't see any way around it right now. Also we
    // should think about splitting this up, but I'm fine -- if a little
    // embarassed about it -- right now. This function swallows errors, and
    // there's a chance we'll end up with no images at the end, but I think we
    // can handle that in the extension's UI.
    pub fn grab_and_store(
        &self,
        client: Client<HttpsConnector<HttpConnector>, Body>,
        db: &Pool<ConnectionManager<PgConnection>>,
        pool: CpuPool,
    ) -> Box<Future<Item = (), Error = ()>> {
        let ad = self.clone();
        let pool_s3 = pool.clone();
        let pool_db = pool.clone();
        let db = db.clone();
        let future = stream::iter_ok(self.image_urls())
            .map(|u| u.unwrap())
            // filter ones we already have in the db and ones we can verify as
            // coming from fb, we don't want to become a malware vector :)
            // currently we redownload images we already have, but ok.
            .filter(|u| {
                info!("testing {:?}", u.host());
                match u.host() {
                    Some(h) => (h == "pp-facebook-ads.s3.amazonaws.com" || h.ends_with("fbcdn.net")),
                    None => false
                }
            })
            // grab image
            .and_then(move |img| {
                info!("getting {:?}", img.path());
                let cloned = img.clone();
                client
                    .get(img)
                    .and_then(|res| {
                        res.body().concat2().and_then(|chunk| Ok((chunk, cloned)))
                    })
                    .map_err(InsertError::Hyper)
            })
            // upload them to s3
            .and_then(move |tuple| {
                let pool = pool_s3.clone();
                // we do this in a worker thread because rusoto isn't on
                // Hyper async yet.
                pool.spawn_fn(move || {
                    if tuple.1.host().unwrap() != "pp-facebook-ads.s3.amazonaws.com" {
                        let credentials = DefaultCredentialsProvider::new()
                            .map_err(InsertError::AWS)?;
                        let tls = default_tls_client()
                            .map_err(InsertError::TLS)?;
                        let client = S3Client::new(tls, credentials, Region::UsEast1);
                        let req = PutObjectRequest {
                            bucket: "pp-facebook-ads".to_string(),
                            key: tuple.1.path().trim_left_matches('/').to_string(),
                            acl: Some("public-read".to_string()),
                            body: Some(tuple.0.to_vec()),
                            ..PutObjectRequest::default()
                        };
                        client.put_object(&req).map_err(InsertError::S3)?;
                    }
                    Ok(tuple.1)
                })
            })
            .collect()
            // save the new urls to the database. the images variable will
            // include only those that we've successfully saved to s3, so we
            // have to do a funky merge here.
            .and_then(move |images| {
                let imgs = images.clone();
                pool_db.spawn_fn(move || {
                    use schema::ads::dsl::*;
                    let update = Images::from_ad(&ad, imgs)?;
                    let connection = db.get().map_err(InsertError::Timeout)?;
                    diesel::update(ads.filter(id.eq(&ad.id)))
                        .set(&update)
                        .execute(&*connection)
                        .map_err(InsertError::DataBase)?;
                    info!("saved {:?}", ad.id);
                    Ok(())
                })
            })
            .map_err(|e| {
                warn!("{:?}", e);
                ()
            });
        Box::new(future)
    }

    pub(self) fn image_urls(&self) -> Vec<Result<Uri, InsertError>> {
        let images = [vec![self.thumbnail.clone()], self.images.clone()];
        images
            .concat()
            .iter()
            .map(|a| {
                let a: Uri = a.parse().map_err(InsertError::Uri)?;
                Ok(a)
            })
            .collect()
    }

    pub fn get_ads_by_lang(
        language: &str,
        conn: &Pool<ConnectionManager<PgConnection>>,
    ) -> Result<Vec<Ad>, InsertError> {
        use schema::ads::dsl::*;
        let connection = conn.get().map_err(InsertError::Timeout)?;
        ads.filter(lang.eq(language))
            .filter(political_probability.gt(0.75))
            .order(created_at.desc())
            .limit(100)
            .load::<Ad>(&*connection)
            .map_err(InsertError::DataBase)
    }
}

#[derive(Insertable)]
#[table_name = "ads"]
pub struct NewAd<'a> {
    id: &'a str,
    html: &'a str,
    political: i32,
    not_political: i32,

    title: String,
    message: String,
    thumbnail: String,

    lang: &'a str,
    images: Vec<String>,
    impressions: i32,

    targeting: Option<String>,
}

impl<'a> NewAd<'a> {
    pub fn new(ad: &'a AdPost, lang: &'a str) -> Result<NewAd<'a>, InsertError> {
        info!("saving {}", ad.id);
        let document = kuchiki::parse_html().one(ad.html.clone());

        let thumb = get_image(&document)?;
        let images = get_images(&document)?;
        let message = get_message(&document)?;
        let title = get_title(&document)?;

        Ok(NewAd {
            id: &ad.id,
            html: &ad.html,
            // we try unwrapping or we chose the false branch in both of these
            // cases to count impressions
            political: if ad.political.unwrap_or(false) { 1 } else { 0 },
            not_political: if !ad.political.unwrap_or(true) { 1 } else { 0 },
            title: title,
            message: message,
            thumbnail: thumb,
            lang: lang,
            images: images,
            impressions: if !ad.political.is_some() { 1 } else { 0 },
            targeting: ad.targeting.clone(),
        })
    }

    pub fn save(&self, pool: &Pool<ConnectionManager<PgConnection>>) -> Result<Ad, InsertError> {
        use schema::ads;
        use schema::ads::dsl::*;
        let connection = pool.get().map_err(InsertError::Timeout)?;

        // increment impressions if this is a background save,
        // otherwise increment political counters
        let ad: Ad = diesel::insert(&self.on_conflict(
            id,
            do_update().set((
                political.eq(political + self.political),
                not_political.eq(
                    not_political +
                        self.not_political,
                ),
                impressions.eq(
                    impressions + self.impressions,
                ),
                updated_at.eq(Utc::now()),
            )),
        )).into(ads::table)
            .get_result(&*connection)
            .map_err(InsertError::DataBase)?;

        if self.targeting.is_some() && !ad.targeting.is_some() {
            diesel::update(ads.filter(id.eq(self.id)))
                .set(targeting.eq(self.targeting.clone()))
                .execute(&*connection)
                .map_err(InsertError::DataBase)?;
        };

        Ok(ad)
    }
}

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

    #[test]
    fn ad_parsing() {
        let ad = include_str!("./html-test.txt");
        let post = AdPost {
            id: "test".to_string(),
            html: ad.to_string(),
            political: None,
            targeting: None,
        };
        let new_ad = NewAd::new(&post, "en-US").unwrap();
        assert!(new_ad.thumbnail.len() > 0);
        assert_eq!(new_ad.images.len(), 2);
        assert!(new_ad.title.len() > 0);

        assert_eq!(
            kuchiki::parse_html().one(new_ad.message).text_contents(),
            kuchiki::parse_html().one("
#Valerian is “the best experience since ‘Avatar.’” See it in 3D and RealD3D theaters this Friday. Get tickets now: ValerianTickets.com
").text_contents()
        );
    }

    #[test]
    fn image_parsing() {
        use chrono::prelude::*;
        let ad = include_str!("./html-test.txt");
        let document = kuchiki::parse_html().one(ad);

        let saved_ad = Ad {
            id: "test".to_string(),
            html: ad.to_string(),
            political: 1,
            not_political: 2,
            title: get_title(&document).unwrap(),
            message: get_message(&document).unwrap(),
            thumbnail: get_image(&document).unwrap(),
            created_at: Utc::now(),
            updated_at: Utc::now(),
            lang: "US".to_string(),
            images: get_images(&document).unwrap(),
            impressions: 1,
            targeting: None,
            political_probability: 0.0,
        };
        let urls = saved_ad
            .image_urls()
            .into_iter()
            .map(|x| x.unwrap())
            .collect();
        let images = Images::from_ad(&saved_ad, urls).unwrap();
        assert!(images.html != saved_ad.html);
        assert!(!images.html.contains("fbcdn"));
        assert!(!images.html.contains("html"));
        assert!(images.images.len() == saved_ad.images.len());
        assert!(images.thumbnail.unwrap() != saved_ad.thumbnail);
    }
}
<|endoftext|>


<commit_before>//! Implements `StdHashTable` and `StdHashTableEntry`.

use libc;
use libc::c_void;
use std::sync::atomic::{AtomicPtr, Ordering};
use std::slice;
use std::isize;
use std::cell::Cell;
use std::cmp::{min, max};
use std::mem;
use value::*;
use depth::*;
use hash_table::*;
use moves::MoveDigest;

/// Implements the `HashTableEntry` trait.
#[derive(Copy, Clone, Debug)]
pub struct StdHashTableEntry {
    value: Value,

    // The transposition table maintains a generation number for each
    // entry, which is used to implement an efficient replacement
    // strategy. This field stores the entry's generation (the highest
    // 6 bits) and the bound type (the lowest 2 bits).
    gen_bound: u8,

    depth: Depth,
    move_digest: MoveDigest,
    static_eval: Value,
}

impl HashTableEntry for StdHashTableEntry {
    #[inline]
    fn new(value: Value,
           bound: BoundType,
           depth: Depth,
           move_digest: MoveDigest)
           -> StdHashTableEntry {
        Self::with_static_eval(value, bound, depth, move_digest, VALUE_UNKNOWN)
    }

    #[inline]
    fn with_static_eval(value: Value,
                        bound: BoundType,
                        depth: Depth,
                        move_digest: MoveDigest,
                        static_eval: Value)
                        -> StdHashTableEntry {
        debug_assert!(value != VALUE_UNKNOWN);
        debug_assert!(bound <= 0b11);
        debug_assert!(DEPTH_MIN <= depth && depth <= DEPTH_MAX);
        StdHashTableEntry {
            value: value,
            gen_bound: bound,
            depth: depth,
            move_digest: move_digest,
            static_eval: static_eval,
        }
    }

    #[inline]
    fn value(&self) -> Value {
        self.value
    }

    #[inline]
    fn bound(&self) -> BoundType {
        self.gen_bound & 0b11
    }

    #[inline]
    fn depth(&self) -> Depth {
        self.depth
    }

    #[inline]
    fn move_digest(&self) -> MoveDigest {
        self.move_digest
    }

    /// Returns the `static_eval` passed to the constructor.
    #[inline]
    fn static_eval(&self) -> Value {
        self.static_eval
    }
}

impl StdHashTableEntry {
    /// Returns the contained data as one `u64` value.
    #[inline]
    fn as_u64(&self) -> u64 {
        unsafe { mem::transmute(*self) }
    }
}

/// Implements the `HashTable` trait.
pub struct StdHashTable {
    /// The current generation number. The lowest 2 bits will always
    /// be zeros.
    generation: Cell<u8>,

    /// The number of buckets in the table.
    bucket_count: usize,

    /// The raw pointer obtained from `libc::calloc`. It will be
    /// passed to `libc::free` before the transposition table is
    /// dropped.
    alloc_ptr: AtomicPtr<c_void>,

    /// The transposition table consists of a vector of buckets. Each
    /// buckets stores 4 records.
    table_ptr: AtomicPtr<[Record; 4]>,
}

impl HashTable for StdHashTable {
    type Entry = StdHashTableEntry;

    fn new(size_mb: Option<usize>) -> StdHashTable {
        let size_mb = size_mb.unwrap_or(16);
        let bucket_size = mem::size_of::<[Record; 4]>();
        let bucket_count = {
            let n = max(1, ((size_mb * 1024 * 1024) / bucket_size) as u64);
            1 << (63 - n.leading_zeros())
        };
        let alloc_ptr;
        let table_ptr = unsafe {
            // Make sure that the first bucket is optimally
            // aligned. This may improve performance when cache line's
            // size is divisible to bucket's size or vice versa.
            alloc_ptr = libc::calloc(bucket_count + 1, bucket_size);
            let mut addr = mem::transmute::<*mut c_void, usize>(alloc_ptr);
            addr += bucket_size;
            addr &= !(bucket_size - 1);
            mem::transmute::<usize, *mut [Record; 4]>(addr)
        };

        StdHashTable {
            generation: Cell::new(0),
            bucket_count: bucket_count,
            alloc_ptr: AtomicPtr::new(alloc_ptr),
            table_ptr: AtomicPtr::new(table_ptr),
        }
    }

    fn new_search(&self) {
        const N: usize = 128;

        loop {
            // Increment `self.generation` (with wrapping).
            self.generation.set(self.generation.get().wrapping_add(0b100));
            debug_assert_eq!(self.generation.get() & 0b11, 0);

            // Count how many staled records from this generation
            // there are among the first `N` buckets.
            let mut staled = 0;
            let mut bucket_iter = self.table().iter();
            for _ in 0..min(N, self.bucket_count) {
                for record in bucket_iter.next().unwrap() {
                    if record.key != 0 && record.generation() == self.generation.get() {
                        staled += 1;
                    }
                }
            }

            if staled < N {
                // Note that we will continue to increment
                // `self.generation` if the staled records from this
                // generation are too many. (This may happen if a very
                // long search was executed long time ago.)
                break;
            }
        }
    }

    fn store(&self, key: u64, mut data: Self::Entry) {
        // `store` and `probe` jointly implement a clever lock-less
        // hashing strategy. Rather than storing two disjoint items,
        // the key is stored XOR-ed with data, while data is stored
        // additionally as usual.

        // Set entry's generation.
        data.gen_bound = self.generation.get() | data.bound();

        // Choose a slot to which to write the data. (Each bucket has
        // 4 slots.)
        let mut bucket = self.bucket_mut(key);
        let mut replace_index = 0;
        let mut replace_score = isize::MAX;
        for (i, record) in bucket.iter_mut().enumerate() {
            // Check if this is an empty slot, or an old record for
            // the same key. If this this is the case we will use this
            // slot for the new record.
            if record.key == 0 || record.key ^ record.data.as_u64() == key {
                if data.move_digest == MoveDigest::invalid() {
                    data.move_digest = record.data.move_digest; // Preserve any existing move.
                }
                replace_index = i;
                break;
            }

            // Calculate the score for this record. If we can not find
            // an empty slot or an old record, the replaced record
            // will be the record with the lowest score.
            let record_score = self.calc_score(record);
            if record_score < replace_score {
                replace_index = i;
                replace_score = record_score;
            }
        }

        // Write the data to the chosen slot.
        bucket[replace_index] = Record {
            key: key ^ data.as_u64(),
            data: data,
        };
    }

    #[inline]
    fn probe(&self, key: u64) -> Option<Self::Entry> {
        // `store` and `probe` jointly implement a clever lock-less
        // hashing strategy. Rather than storing two disjoint items,
        // the key is stored XOR-ed with data, while data is stored
        // additionally as usual.

        let bucket = self.bucket_mut(key);
        for record in bucket.iter_mut() {
            if record.key ^ record.data.as_u64() == key {
                // If `key` and `data` were written simultaneously by
                // different search instances with different keys,
                // this will yield in a mismatch of the above
                // comparison (except for the rare and inherent key
                // collisions).
                record.set_generation(self.generation.get());
                return Some(record.data);
            }
        }
        None
    }

    fn clear(&self) {
        for bucket in self.table_mut() {
            for record in bucket.iter_mut() {
                *record = Default::default();
            }
        }
        self.generation.set(0);
    }
}

impl StdHashTable {
    /// A helper method for `store`. It implements the record
    /// replacement strategy.
    #[inline]
    fn calc_score(&self, record: &Record) -> isize {
        // Here we try to return higher values for the records that
        // are move likely to save CPU work in the future:

        // Positions from the current generation are always scored
        // higher than positions from older generations.
        (if record.generation() == self.generation.get() {
            DEPTH_MAX as isize + 2
        } else {
            0
        })

        // Positions with higher search depths are scored higher.
        + record.data.depth() as isize

        // Positions with exact evaluations are given slight advantage.
        + (if record.data.bound() == BOUND_EXACT {
            1
        } else {
            0
        })
    }

    /// A helper method for `probe` and `store`. It returns the bucket
    /// for a given key.
    #[inline]
    fn bucket_mut(&self, key: u64) -> &mut [Record; 4] {
        unsafe {
            let index = (key & (self.bucket_count - 1) as u64) as usize;
            self.table_mut().get_unchecked_mut(index)
        }
    }

    #[inline]
    fn table(&self) -> &[[Record; 4]] {
        unsafe { slice::from_raw_parts(self.table_ptr.load(Ordering::Relaxed), self.bucket_count) }
    }

    #[inline]
    fn table_mut(&self) -> &mut [[Record; 4]] {
        unsafe {
            slice::from_raw_parts_mut(self.table_ptr.load(Ordering::Relaxed), self.bucket_count)
        }
    }
}

impl Drop for StdHashTable {
    fn drop(&mut self) {
        unsafe {
            libc::free(self.alloc_ptr.load(Ordering::Relaxed));
        }
    }
}

unsafe impl Sync for StdHashTable {}

/// Represents a record in the transposition table.
///
/// It consists of 16 bytes, and is laid out the following way:
///
/// * key         64 bit
/// * move_digest 16 bit
/// * value       16 bit
/// * eval value  16 bit
/// * depth        8 bit
/// * generation   6 bit
/// * bound type   2 bit
#[derive(Copy, Clone)]
struct Record {
    key: u64,
    data: StdHashTableEntry,
}

impl Default for Record {
    fn default() -> Record {
        Record {
            key: 0,
            data: unsafe { mem::transmute(0u64) },
        }
    }
}

impl Record {
    #[inline]
    fn generation(&self) -> u8 {
        self.data.gen_bound & 0b11111100
    }

    #[inline]
    fn set_generation(&mut self, generation: u8) {
        debug_assert_eq!(generation & 0b11, 0);

        // Since the `key` is saved XOR-ed with the data, when we
        // change the data, we have to change the stored `key` as
        // well.
        let old_data_as_u64 = self.data.as_u64();
        self.data.gen_bound = generation | self.data.bound();
        self.key ^= old_data_as_u64 ^ self.data.as_u64();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use super::Record;
    use std;
    use depth::*;
    use hash_table::*;
    use moves::*;

    #[test]
    fn bucket_size() {
        assert_eq!(std::mem::size_of::<[Record; 4]>(), 64);
        assert_eq!(std::mem::size_of::<Record>(), 16);
    }

    #[test]
    fn store_and_probe() {
        let tt = StdHashTable::new(None);
        assert!(tt.probe(1).is_none());
        let data = StdHashTableEntry::new(0, 0, 50, MoveDigest::invalid());
        assert_eq!(data.depth(), 50);
        assert_eq!(data.move_digest(), MoveDigest::invalid());
        tt.store(1, data);
        assert_eq!(tt.probe(1).unwrap().depth(), 50);
        tt.store(1, StdHashTableEntry::new(0, 0, 50, MoveDigest::invalid()));
        assert_eq!(tt.probe(1).unwrap().depth(), 50);
        assert_eq!(tt.probe(1).unwrap().move_digest(), MoveDigest::invalid());
        let digest = MoveDigest::invalid();
        for i in 2..50 {
            tt.store(i, StdHashTableEntry::new(i as i16, 0, i as Depth, digest));
        }
        assert_eq!(tt.probe(1).unwrap().depth(), 50);
        assert_eq!(tt.probe(49).unwrap().depth(), 49);
        assert_eq!(tt.probe(48).unwrap().depth(), 48);
        assert_eq!(tt.probe(47).unwrap().depth(), 47);
        tt.clear();
        assert!(tt.probe(1).is_none());
        tt.store(1, data);
        tt.new_search();
        tt.probe(1);
        assert!(tt.probe(1).is_some());
    }

    #[test]
    fn new_search() {
        let tt = StdHashTable::new(None);
        assert_eq!(tt.generation.get(), 0 << 2);
        tt.new_search();
        assert_eq!(tt.generation.get(), 1 << 2);
        for _ in 0..64 {
            tt.new_search();
        }
        assert_eq!(tt.generation.get(), 1 << 2);
    }
}
<commit_msg>add Bucket type<commit_after>//! Implements `StdHashTable` and `StdHashTableEntry`.

use libc;
use libc::c_void;
use std::sync::atomic::{AtomicPtr, Ordering};
use std::slice;
use std::isize;
use std::cell::Cell;
use std::cmp::{min, max};
use std::mem;
use value::*;
use depth::*;
use hash_table::*;
use moves::MoveDigest;

/// Implements the `HashTableEntry` trait.
#[derive(Copy, Clone, Debug)]
pub struct StdHashTableEntry {
    value: Value,

    // The transposition table maintains a generation number for each
    // entry, which is used to implement an efficient replacement
    // strategy. This field stores the entry's generation (the highest
    // 6 bits) and the bound type (the lowest 2 bits).
    gen_bound: u8,

    depth: Depth,
    move_digest: MoveDigest,
    static_eval: Value,
}

impl HashTableEntry for StdHashTableEntry {
    #[inline]
    fn new(value: Value,
           bound: BoundType,
           depth: Depth,
           move_digest: MoveDigest)
           -> StdHashTableEntry {
        Self::with_static_eval(value, bound, depth, move_digest, VALUE_UNKNOWN)
    }

    #[inline]
    fn with_static_eval(value: Value,
                        bound: BoundType,
                        depth: Depth,
                        move_digest: MoveDigest,
                        static_eval: Value)
                        -> StdHashTableEntry {
        debug_assert!(value != VALUE_UNKNOWN);
        debug_assert!(bound <= 0b11);
        debug_assert!(DEPTH_MIN <= depth && depth <= DEPTH_MAX);
        StdHashTableEntry {
            value: value,
            gen_bound: bound,
            depth: depth,
            move_digest: move_digest,
            static_eval: static_eval,
        }
    }

    #[inline]
    fn value(&self) -> Value {
        self.value
    }

    #[inline]
    fn bound(&self) -> BoundType {
        self.gen_bound & 0b11
    }

    #[inline]
    fn depth(&self) -> Depth {
        self.depth
    }

    #[inline]
    fn move_digest(&self) -> MoveDigest {
        self.move_digest
    }

    /// Returns the `static_eval` passed to the constructor.
    #[inline]
    fn static_eval(&self) -> Value {
        self.static_eval
    }
}

impl StdHashTableEntry {
    /// Returns the contained data as one `u64` value.
    #[inline]
    fn as_u64(&self) -> u64 {
        unsafe { mem::transmute(*self) }
    }
}

/// Implements the `HashTable` trait.
pub struct StdHashTable {
    /// The current generation number. The lowest 2 bits will always
    /// be zeros.
    generation: Cell<u8>,

    /// The number of buckets in the table.
    bucket_count: usize,

    /// The raw pointer obtained from `libc::calloc`. It will be
    /// passed to `libc::free` before the transposition table is
    /// dropped.
    alloc_ptr: AtomicPtr<c_void>,

    /// The transposition table consists of a vector of buckets. Each
    /// buckets stores 4 records.
    table_ptr: AtomicPtr<Bucket>,
}

impl HashTable for StdHashTable {
    type Entry = StdHashTableEntry;

    fn new(size_mb: Option<usize>) -> StdHashTable {
        let size_mb = size_mb.unwrap_or(16);
        let bucket_size = mem::size_of::<Bucket>();
        let bucket_count = {
            let n = max(1, ((size_mb * 1024 * 1024) / bucket_size) as u64);
            1 << (63 - n.leading_zeros())
        };
        let alloc_ptr;
        let table_ptr = unsafe {
            // Make sure that the first bucket is optimally
            // aligned. This may improve performance when cache line's
            // size is divisible to bucket's size or vice versa.
            alloc_ptr = libc::calloc(bucket_count + 1, bucket_size);
            let mut addr = mem::transmute::<*mut c_void, usize>(alloc_ptr);
            addr += bucket_size;
            addr &= !(bucket_size - 1);
            mem::transmute::<usize, *mut Bucket>(addr)
        };

        StdHashTable {
            generation: Cell::new(0),
            bucket_count: bucket_count,
            alloc_ptr: AtomicPtr::new(alloc_ptr),
            table_ptr: AtomicPtr::new(table_ptr),
        }
    }

    fn new_search(&self) {
        const N: usize = 128;

        loop {
            // Increment `self.generation` (with wrapping).
            self.generation.set(self.generation.get().wrapping_add(0b100));
            debug_assert_eq!(self.generation.get() & 0b11, 0);

            // Count how many staled records from this generation
            // there are among the first `N` buckets.
            let mut staled = 0;
            let mut bucket_iter = self.table().iter();
            for _ in 0..min(N, self.bucket_count) {
                for record in bucket_iter.next().unwrap() {
                    if record.key != 0 && record.generation() == self.generation.get() {
                        staled += 1;
                    }
                }
            }

            if staled < N {
                // Note that we will continue to increment
                // `self.generation` if the staled records from this
                // generation are too many. (This may happen if a very
                // long search was executed long time ago.)
                break;
            }
        }
    }

    fn store(&self, key: u64, mut data: Self::Entry) {
        // `store` and `probe` jointly implement a clever lock-less
        // hashing strategy. Rather than storing two disjoint items,
        // the key is stored XOR-ed with data, while data is stored
        // additionally as usual.

        // Set entry's generation.
        data.gen_bound = self.generation.get() | data.bound();

        // Choose a slot to which to write the data. (Each bucket can
        // have several slots.)
        let mut bucket = self.bucket_mut(key);
        let mut replace_index = 0;
        let mut replace_score = isize::MAX;
        for (i, record) in bucket.iter_mut().enumerate() {
            // Check if this is an empty slot, or an old record for
            // the same key. If this this is the case we will use this
            // slot for the new record.
            if record.key == 0 || record.key ^ record.data.as_u64() == key {
                if data.move_digest == MoveDigest::invalid() {
                    data.move_digest = record.data.move_digest; // Preserve any existing move.
                }
                replace_index = i;
                break;
            }

            // Calculate the score for this record. If we can not find
            // an empty slot or an old record, the replaced record
            // will be the record with the lowest score.
            let record_score = self.calc_score(record);
            if record_score < replace_score {
                replace_index = i;
                replace_score = record_score;
            }
        }

        // Write the data to the chosen slot.
        bucket[replace_index] = Record {
            key: key ^ data.as_u64(),
            data: data,
        };
    }

    #[inline]
    fn probe(&self, key: u64) -> Option<Self::Entry> {
        // `store` and `probe` jointly implement a clever lock-less
        // hashing strategy. Rather than storing two disjoint items,
        // the key is stored XOR-ed with data, while data is stored
        // additionally as usual.

        let bucket = self.bucket_mut(key);
        for record in bucket.iter_mut() {
            if record.key ^ record.data.as_u64() == key {
                // If `key` and `data` were written simultaneously by
                // different search instances with different keys,
                // this will yield in a mismatch of the above
                // comparison (except for the rare and inherent key
                // collisions).
                record.set_generation(self.generation.get());
                return Some(record.data);
            }
        }
        None
    }

    fn clear(&self) {
        for bucket in self.table_mut() {
            for record in bucket.iter_mut() {
                *record = Default::default();
            }
        }
        self.generation.set(0);
    }
}

impl StdHashTable {
    /// A helper method for `store`. It implements the record
    /// replacement strategy.
    #[inline]
    fn calc_score(&self, record: &Record) -> isize {
        // Here we try to return higher values for the records that
        // are move likely to save CPU work in the future:

        // Positions from the current generation are always scored
        // higher than positions from older generations.
        (if record.generation() == self.generation.get() {
            DEPTH_MAX as isize + 2
        } else {
            0
        })

        // Positions with higher search depths are scored higher.
        + record.data.depth() as isize

        // Positions with exact evaluations are given slight advantage.
        + (if record.data.bound() == BOUND_EXACT {
            1
        } else {
            0
        })
    }

    /// A helper method for `probe` and `store`. It returns the bucket
    /// for a given key.
    #[inline]
    fn bucket_mut(&self, key: u64) -> &mut Bucket {
        unsafe {
            let index = (key & (self.bucket_count - 1) as u64) as usize;
            self.table_mut().get_unchecked_mut(index)
        }
    }

    #[inline]
    fn table(&self) -> &[Bucket] {
        unsafe { slice::from_raw_parts(self.table_ptr.load(Ordering::Relaxed), self.bucket_count) }
    }

    #[inline]
    fn table_mut(&self) -> &mut [Bucket] {
        unsafe {
            slice::from_raw_parts_mut(self.table_ptr.load(Ordering::Relaxed), self.bucket_count)
        }
    }
}

impl Drop for StdHashTable {
    fn drop(&mut self) {
        unsafe {
            libc::free(self.alloc_ptr.load(Ordering::Relaxed));
        }
    }
}

unsafe impl Sync for StdHashTable {}

/// Represents a record in the transposition table.
///
/// It consists of 16 bytes, and is laid out the following way:
///
/// * key         64 bit
/// * move_digest 16 bit
/// * value       16 bit
/// * eval value  16 bit
/// * depth        8 bit
/// * generation   6 bit
/// * bound type   2 bit
#[derive(Copy, Clone)]
struct Record {
    key: u64,
    data: StdHashTableEntry,
}

impl Default for Record {
    fn default() -> Record {
        Record {
            key: 0,
            data: unsafe { mem::transmute(0u64) },
        }
    }
}

impl Record {
    #[inline]
    fn generation(&self) -> u8 {
        self.data.gen_bound & 0b11111100
    }

    #[inline]
    fn set_generation(&mut self, generation: u8) {
        debug_assert_eq!(generation & 0b11, 0);

        // Since the `key` is saved XOR-ed with the data, when we
        // change the data, we have to change the stored `key` as
        // well.
        let old_data_as_u64 = self.data.as_u64();
        self.data.gen_bound = generation | self.data.bound();
        self.key ^= old_data_as_u64 ^ self.data.as_u64();
    }
}

type Bucket = [Record; 4];

#[cfg(test)]
mod tests {
    use super::*;
    use super::Bucket;
    use std;
    use depth::*;
    use hash_table::*;
    use moves::*;

    #[test]
    fn bucket_size() {
        let n = std::mem::size_of::<Bucket>();
        assert!(n % 64 == 0 || 64 % n == 0);
    }

    #[test]
    fn store_and_probe() {
        let tt = StdHashTable::new(None);
        assert!(tt.probe(1).is_none());
        let data = StdHashTableEntry::new(0, 0, 50, MoveDigest::invalid());
        assert_eq!(data.depth(), 50);
        assert_eq!(data.move_digest(), MoveDigest::invalid());
        tt.store(1, data);
        assert_eq!(tt.probe(1).unwrap().depth(), 50);
        tt.store(1, StdHashTableEntry::new(0, 0, 50, MoveDigest::invalid()));
        assert_eq!(tt.probe(1).unwrap().depth(), 50);
        assert_eq!(tt.probe(1).unwrap().move_digest(), MoveDigest::invalid());
        let digest = MoveDigest::invalid();
        for i in 2..50 {
            tt.store(i, StdHashTableEntry::new(i as i16, 0, i as Depth, digest));
        }
        assert_eq!(tt.probe(1).unwrap().depth(), 50);
        assert_eq!(tt.probe(49).unwrap().depth(), 49);
        assert_eq!(tt.probe(48).unwrap().depth(), 48);
        assert_eq!(tt.probe(47).unwrap().depth(), 47);
        tt.clear();
        assert!(tt.probe(1).is_none());
        tt.store(1, data);
        tt.new_search();
        tt.probe(1);
        assert!(tt.probe(1).is_some());
    }

    #[test]
    fn new_search() {
        let tt = StdHashTable::new(None);
        assert_eq!(tt.generation.get(), 0 << 2);
        tt.new_search();
        assert_eq!(tt.generation.get(), 1 << 2);
        for _ in 0..64 {
            tt.new_search();
        }
        assert_eq!(tt.generation.get(), 1 << 2);
    }
}
<|endoftext|>


<commit_before>#![feature(doc_cfg)]
#![feature(target_feature, cfg_target_feature)]

// @has doc_cfg/struct.Portable.html
// @!has - '//*[@id="main"]/*[@class="item-info"]/*[@class="stab portability"]' ''
// @has - '//*[@id="method.unix_and_arm_only_function"]' 'fn unix_and_arm_only_function()'
// @has - '//*[@class="stab portability"]' 'This is supported on Unix and ARM only.'
pub struct Portable;

// @has doc_cfg/unix_only/index.html \
//  '//*[@id="main"]/*[@class="item-info"]/*[@class="stab portability"]' \
//  'This is supported on Unix only.'
// @matches - '//*[@class="module-item"]//*[@class="stab portability"]' '\AARM\Z'
// @count - '//*[@class="stab portability"]' 2
#[doc(cfg(unix))]
pub mod unix_only {
    // @has doc_cfg/unix_only/fn.unix_only_function.html \
    //  '//*[@id="main"]/*[@class="item-info"]/*[@class="stab portability"]' \
    //  'This is supported on Unix only.'
    // @count - '//*[@class="stab portability"]' 1
    pub fn unix_only_function() {
        content::should::be::irrelevant();
    }

    // @has doc_cfg/unix_only/trait.ArmOnly.html \
    //  '//*[@id="main"]/*[@class="item-info"]/*[@class="stab portability"]' \
    //  'This is supported on Unix and ARM only.'
    // @count - '//*[@class="stab portability"]' 1
    #[doc(cfg(target_arch = "arm"))]
    pub trait ArmOnly {
        fn unix_and_arm_only_function();
    }

    #[doc(cfg(target_arch = "arm"))]
    impl ArmOnly for super::Portable {
        fn unix_and_arm_only_function() {}
    }
}

// tagging a function with `#[target_feature]` creates a doc(cfg(target_feature)) node for that
// item as well

// the portability header is different on the module view versus the full view
// @has doc_cfg/index.html
// @matches - '//*[@class="module-item"]//*[@class="stab portability"]' '\Aavx\Z'

// @has doc_cfg/fn.uses_target_feature.html
// @has - '//*[@id="main"]/*[@class="item-info"]/*[@class="stab portability"]' \
//        'This is supported with target feature avx only.'
#[target_feature(enable = "avx")]
pub unsafe fn uses_target_feature() {
    content::should::be::irrelevant();
}

// @has doc_cfg/fn.uses_cfg_target_feature.html
// @has - '//*[@id="main"]/*[@class="item-info"]/*[@class="stab portability"]' \
//        'This is supported with target feature avx only.'
#[doc(cfg(target_feature = "avx"))]
pub fn uses_cfg_target_feature() {
    uses_target_feature();
}
<commit_msg>Add a test in src/test/rustdoc/doc-cfg.rs<commit_after>#![feature(doc_cfg)]
#![feature(target_feature, cfg_target_feature)]

// @has doc_cfg/struct.Portable.html
// @!has - '//*[@id="main"]/*[@class="item-info"]/*[@class="stab portability"]' ''
// @has - '//*[@id="method.unix_and_arm_only_function"]' 'fn unix_and_arm_only_function()'
// @has - '//*[@class="stab portability"]' 'This is supported on Unix and ARM only.'
// @has - '//*[@id="method.wasi_and_wasm32_only_function"]' 'fn wasi_and_wasm32_only_function()'
// @has - '//*[@class="stab portability"]' 'This is supported on WASI and WebAssembly only.'
pub struct Portable;

// @has doc_cfg/unix_only/index.html \
//  '//*[@id="main"]/*[@class="item-info"]/*[@class="stab portability"]' \
//  'This is supported on Unix only.'
// @matches - '//*[@class="module-item"]//*[@class="stab portability"]' '\AARM\Z'
// @count - '//*[@class="stab portability"]' 2
#[doc(cfg(unix))]
pub mod unix_only {
    // @has doc_cfg/unix_only/fn.unix_only_function.html \
    //  '//*[@id="main"]/*[@class="item-info"]/*[@class="stab portability"]' \
    //  'This is supported on Unix only.'
    // @count - '//*[@class="stab portability"]' 1
    pub fn unix_only_function() {
        content::should::be::irrelevant();
    }

    // @has doc_cfg/unix_only/trait.ArmOnly.html \
    //  '//*[@id="main"]/*[@class="item-info"]/*[@class="stab portability"]' \
    //  'This is supported on Unix and ARM only.'
    // @count - '//*[@class="stab portability"]' 1
    #[doc(cfg(target_arch = "arm"))]
    pub trait ArmOnly {
        fn unix_and_arm_only_function();
    }

    #[doc(cfg(target_arch = "arm"))]
    impl ArmOnly for super::Portable {
        fn unix_and_arm_only_function() {}
    }
}

// @has doc_cfg/wasi_only/index.html \
//  '//*[@id="main"]/*[@class="item-info"]/*[@class="stab portability"]' \
//  'This is supported on WASI only.'
// @matches - '//*[@class="module-item"]//*[@class="stab portability"]' '\AWebAssembly\Z'
// @count - '//*[@class="stab portability"]' 2
#[doc(cfg(target_os = "wasi"))]
pub mod wasi_only {
    // @has doc_cfg/wasi_only/fn.wasi_only_function.html \
    //  '//*[@id="main"]/*[@class="item-info"]/*[@class="stab portability"]' \
    //  'This is supported on WASI only.'
    // @count - '//*[@class="stab portability"]' 1
    pub fn wasi_only_function() {
        content::should::be::irrelevant();
    }

    // @has doc_cfg/wasi_only/trait.Wasm32Only.html \
    //  '//*[@id="main"]/*[@class="item-info"]/*[@class="stab portability"]' \
    //  'This is supported on WASI and WebAssembly only.'
    // @count - '//*[@class="stab portability"]' 1
    #[doc(cfg(target_arch = "wasm32"))]
    pub trait Wasm32Only {
        fn wasi_and_wasm32_only_function();
    }

    #[doc(cfg(target_arch = "wasm32"))]
    impl Wasm32Only for super::Portable {
        fn wasi_and_wasm32_only_function() {}
    }
}

// tagging a function with `#[target_feature]` creates a doc(cfg(target_feature)) node for that
// item as well

// the portability header is different on the module view versus the full view
// @has doc_cfg/index.html
// @matches - '//*[@class="module-item"]//*[@class="stab portability"]' '\Aavx\Z'

// @has doc_cfg/fn.uses_target_feature.html
// @has - '//*[@id="main"]/*[@class="item-info"]/*[@class="stab portability"]' \
//        'This is supported with target feature avx only.'
#[target_feature(enable = "avx")]
pub unsafe fn uses_target_feature() {
    content::should::be::irrelevant();
}

// @has doc_cfg/fn.uses_cfg_target_feature.html
// @has - '//*[@id="main"]/*[@class="item-info"]/*[@class="stab portability"]' \
//        'This is supported with target feature avx only.'
#[doc(cfg(target_feature = "avx"))]
pub fn uses_cfg_target_feature() {
    uses_target_feature();
}
<|endoftext|>


<commit_before>// * This file is part of the uutils coreutils package.
// *
// * (c) 2014 T. Jameson Little <t.jameson.little@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.

extern crate rand;

#[macro_use]
extern crate uucore;

use std::collections::HashMap;
use std::fmt;
use std::io::{stdin, BufRead};
use std::ops;

mod miller_rabin;
mod numeric;
mod rho;
mod table;

static SYNTAX: &str = "[OPTION] [NUMBER]...";
static SUMMARY: &str = "Print the prime factors of the given number(s).
 If none are specified, read from standard input.";
static LONG_HELP: &str = "";

struct Factors {
    f: HashMap<u64, u8>,
}

impl Factors {
    fn new() -> Factors {
        Factors { f: HashMap::new() }
    }

    fn add(&mut self, prime: u64, exp: u8) {
        debug_assert!(exp > 0);
        let n = *self.f.get(&prime).unwrap_or(&0);
        self.f.insert(prime, exp + n);
    }

    fn push(&mut self, prime: u64) {
        self.add(prime, 1)
    }

    #[cfg(test)]
    fn product(&self) -> u64 {
        self.f
            .iter()
            .fold(1, |acc, (p, exp)| acc * p.pow(*exp as u32))
    }
}

impl ops::MulAssign<Factors> for Factors {
    fn mul_assign(&mut self, other: Factors) {
        for (prime, exp) in &other.f {
            self.add(*prime, *exp);
        }
    }
}

impl fmt::Display for Factors {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // TODO: Use a representation with efficient in-order iteration
        let mut primes: Vec<&u64> = self.f.keys().collect();
        primes.sort();

        for p in primes {
            for _ in 0..self.f[&p] {
                write!(f, " {}", p)?
            }
        }

        Ok(())
    }
}

fn factor(mut n: u64) -> Factors {
    let mut factors = Factors::new();

    if n < 2 {
        factors.push(n);
        return factors;
    }

    let z = n.trailing_zeros();
    if z > 0 {
        factors.add(2, z as u8);
        n >>= z;
    }

    if n == 1 {
        return factors;
    }

    let (f, n) = table::factor(n);
    factors *= f;

    if n >= table::NEXT_PRIME {
        factors *= rho::factor(n);
    }

    factors
}

fn print_factors(num: u64) {
    print!("{}:{}", num, factor(num));
    println!();
}

fn print_factors_str(num_str: &str) {
    if let Err(e) = num_str.parse::<u64>().and_then(|x| {
        print_factors(x);
        Ok(())
    }) {
        show_warning!("{}: {}", num_str, e);
    }
}

pub fn uumain(args: impl uucore::Args) -> i32 {
    let matches = app!(SYNTAX, SUMMARY, LONG_HELP).parse(args.collect_str());

    if matches.free.is_empty() {
        let stdin = stdin();
        for line in stdin.lock().lines() {
            for number in line.unwrap().split_whitespace() {
                print_factors_str(number);
            }
        }
    } else {
        for num_str in &matches.free {
            print_factors_str(num_str);
        }
    }
    0
}

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

    #[test]
    fn factor_recombines_small() {
        assert!((1..10_000)
            .map(|i| 2 * i + 1)
            .all(|i| factor(i).product() == i));
    }

    #[test]
    fn factor_recombines_overflowing() {
        assert!((0..250)
            .map(|i| 2 * i + 2u64.pow(32) + 1)
            .all(|i| factor(i).product() == i));
    }
}
<commit_msg>factor::Factors: Use a tree-based map internally<commit_after>// * This file is part of the uutils coreutils package.
// *
// * (c) 2014 T. Jameson Little <t.jameson.little@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.

extern crate rand;

#[macro_use]
extern crate uucore;

use std::collections::BTreeMap;
use std::fmt;
use std::io::{stdin, BufRead};
use std::ops;

mod miller_rabin;
mod numeric;
mod rho;
mod table;

static SYNTAX: &str = "[OPTION] [NUMBER]...";
static SUMMARY: &str = "Print the prime factors of the given number(s).
 If none are specified, read from standard input.";
static LONG_HELP: &str = "";

struct Factors {
    f: BTreeMap<u64, u8>,
}

impl Factors {
    fn new() -> Factors {
        Factors { f: BTreeMap::new() }
    }

    fn add(&mut self, prime: u64, exp: u8) {
        debug_assert!(exp > 0);
        let n = *self.f.get(&prime).unwrap_or(&0);
        self.f.insert(prime, exp + n);
    }

    fn push(&mut self, prime: u64) {
        self.add(prime, 1)
    }

    #[cfg(test)]
    fn product(&self) -> u64 {
        self.f
            .iter()
            .fold(1, |acc, (p, exp)| acc * p.pow(*exp as u32))
    }
}

impl ops::MulAssign<Factors> for Factors {
    fn mul_assign(&mut self, other: Factors) {
        for (prime, exp) in &other.f {
            self.add(*prime, *exp);
        }
    }
}

impl fmt::Display for Factors {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for (p, exp) in self.f.iter() {
            for _ in 0..*exp {
                write!(f, " {}", p)?
            }
        }

        Ok(())
    }
}

fn factor(mut n: u64) -> Factors {
    let mut factors = Factors::new();

    if n < 2 {
        factors.push(n);
        return factors;
    }

    let z = n.trailing_zeros();
    if z > 0 {
        factors.add(2, z as u8);
        n >>= z;
    }

    if n == 1 {
        return factors;
    }

    let (f, n) = table::factor(n);
    factors *= f;

    if n >= table::NEXT_PRIME {
        factors *= rho::factor(n);
    }

    factors
}

fn print_factors(num: u64) {
    print!("{}:{}", num, factor(num));
    println!();
}

fn print_factors_str(num_str: &str) {
    if let Err(e) = num_str.parse::<u64>().and_then(|x| {
        print_factors(x);
        Ok(())
    }) {
        show_warning!("{}: {}", num_str, e);
    }
}

pub fn uumain(args: impl uucore::Args) -> i32 {
    let matches = app!(SYNTAX, SUMMARY, LONG_HELP).parse(args.collect_str());

    if matches.free.is_empty() {
        let stdin = stdin();
        for line in stdin.lock().lines() {
            for number in line.unwrap().split_whitespace() {
                print_factors_str(number);
            }
        }
    } else {
        for num_str in &matches.free {
            print_factors_str(num_str);
        }
    }
    0
}

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

    #[test]
    fn factor_recombines_small() {
        assert!((1..10_000)
            .map(|i| 2 * i + 1)
            .all(|i| factor(i).product() == i));
    }

    #[test]
    fn factor_recombines_overflowing() {
        assert!((0..250)
            .map(|i| 2 * i + 2u64.pow(32) + 1)
            .all(|i| factor(i).product() == i));
    }
}
<|endoftext|>