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

#[cfg(feature = "model")]
use builder::{CreateMessage, GetMessages};
#[cfg(feature = "model")]
use http::AttachmentType;
#[cfg(feature = "model")]
use internal::RwLockExt;

/// A Direct Message text channel with another user.
#[derive(Clone, Debug, Deserialize)]
pub struct PrivateChannel {
    /// The unique Id of the private channel.
    ///
    /// Can be used to calculate the first message's creation date.
    pub id: ChannelId,
    /// The Id of the last message sent.
    pub last_message_id: Option<MessageId>,
    /// Timestamp of the last time a [`Message`] was pinned.
    ///
    /// [`Message`]: struct.Message.html
    pub last_pin_timestamp: Option<DateTime<FixedOffset>>,
    /// Indicator of the type of channel this is.
    ///
    /// This should always be [`ChannelType::Private`].
    ///
    /// [`ChannelType::Private`]: enum.ChannelType.html#variant.Private
    #[serde(rename = "type")]
    pub kind: ChannelType,
    /// The recipient to the private channel.
    #[serde(deserialize_with = "deserialize_single_recipient", rename = "recipients")]
    pub recipient: Arc<RwLock<User>>,
}

#[cfg(feature = "model")]
impl PrivateChannel {
    /// Broadcasts that the current user is typing to the recipient.
    pub fn broadcast_typing(&self) -> Result<()> { self.id.broadcast_typing() }

    /// React to a [`Message`] with a custom [`Emoji`] or unicode character.
    ///
    /// [`Message::react`] may be a more suited method of reacting in most
    /// cases.
    ///
    /// Requires the [Add Reactions] permission, _if_ the current user is the
    /// first user to perform a react with a certain emoji.
    ///
    /// [`Emoji`]: struct.Emoji.html
    /// [`Message`]: struct.Message.html
    /// [`Message::react`]: struct.Message.html#method.react
    /// [Add Reactions]: permissions/constant.ADD_REACTIONS.html
    pub fn create_reaction<M, R>(&self, message_id: M, reaction_type: R) -> Result<()>
        where M: Into<MessageId>, R: Into<ReactionType> {
        self.id.create_reaction(message_id, reaction_type)
    }

    /// Deletes the channel. This does not delete the contents of the channel,
    /// and is equivalent to closing a private channel on the client, which can
    /// be re-opened.
    #[inline]
    pub fn delete(&self) -> Result<Channel> { self.id.delete() }

    /// Deletes all messages by Ids from the given vector in the channel.
    ///
    /// Refer to [`Channel::delete_messages`] for more information.
    ///
    /// Requires the [Manage Messages] permission.
    ///
    /// **Note**: This uses bulk delete endpoint which is not available
    /// for user accounts.
    ///
    /// **Note**: Messages that are older than 2 weeks can't be deleted using
    /// this method.
    ///
    /// [`Channel::delete_messages`]: enum.Channel.html#method.delete_messages
    /// [Manage Messages]: permissions/constant.MANAGE_MESSAGES.html
    #[inline]
    pub fn delete_messages<'a, It: IntoIterator<Item=&'a MessageId>>(&self, message_ids: It) -> Result<()> {
        self.id.delete_messages(message_ids)
    }

    /// Deletes all permission overrides in the channel from a member
    /// or role.
    ///
    /// **Note**: Requires the [Manage Channel] permission.
    ///
    /// [Manage Channel]: permissions/constant.MANAGE_CHANNELS.html
    #[inline]
    pub fn delete_permission(&self, permission_type: PermissionOverwriteType) -> Result<()> {
        self.id.delete_permission(permission_type)
    }

    /// Deletes the given [`Reaction`] from the channel.
    ///
    /// **Note**: Requires the [Manage Messages] permission, _if_ the current
    /// user did not perform the reaction.
    ///
    /// [`Reaction`]: struct.Reaction.html
    /// [Manage Messages]: permissions/constant.MANAGE_MESSAGES.html
    #[inline]
    pub fn delete_reaction<M, R>(&self,
                                 message_id: M,
                                 user_id: Option<UserId>,
                                 reaction_type: R)
                                 -> Result<()>
        where M: Into<MessageId>, R: Into<ReactionType> {
        self.id.delete_reaction(message_id, user_id, reaction_type)
    }

    /// Edits a [`Message`] in the channel given its Id.
    ///
    /// Message editing preserves all unchanged message data.
    ///
    /// Refer to the documentation for [`CreateMessage`] for more information
    /// regarding message restrictions and requirements.
    ///
    /// **Note**: Requires that the current user be the author of the message.
    ///
    /// # Errors
    ///
    /// Returns a [`ModelError::MessageTooLong`] if the content of the message
    /// is over the [`the limit`], containing the number of unicode code points
    /// over the limit.
    ///
    /// [`ModelError::MessageTooLong`]: enum.ModelError.html#variant.MessageTooLong
    /// [`CreateMessage`]: ../builder/struct.CreateMessage.html
    /// [`Message`]: struct.Message.html
    /// [`the limit`]: ../builder/struct.CreateMessage.html#method.content
    #[inline]
    pub fn edit_message<F, M>(&self, message_id: M, f: F) -> Result<Message>
        where F: FnOnce(CreateMessage) -> CreateMessage, M: Into<MessageId> {
        self.id.edit_message(message_id, f)
    }

    /// Determines if the channel is NSFW.
    ///
    /// Refer to [`utils::is_nsfw`] for more details.
    ///
    /// **Note**: This method is for consistency. This will always return
    /// `false`, due to DMs not being considered NSFW.
    ///
    /// [`utils::is_nsfw`]: ../utils/fn.is_nsfw.html
    #[inline]
    pub fn is_nsfw(&self) -> bool { false }

    /// Gets a message from the channel.
    ///
    /// Requires the [Read Message History] permission.
    ///
    /// [Read Message History]: permissions/constant.READ_MESSAGE_HISTORY.html
    #[inline]
    pub fn message<M: Into<MessageId>>(&self, message_id: M) -> Result<Message> {
        self.id.message(message_id)
    }

    /// Gets messages from the channel.
    ///
    /// Refer to [`Channel::messages`] for more information.
    ///
    /// Requires the [Read Message History] permission.
    ///
    /// [`Channel::messages`]: enum.Channel.html#method.messages
    /// [Read Message History]: permissions/constant.READ_MESSAGE_HISTORY.html
    #[inline]
    pub fn messages<F>(&self, f: F) -> Result<Vec<Message>>
        where F: FnOnce(GetMessages) -> GetMessages {
        self.id.messages(f)
    }

    /// Returns "DM with $username#discriminator".
    pub fn name(&self) -> String { format!("DM with {}", self.recipient.with(|r| r.tag())) }

    /// Gets the list of [`User`]s who have reacted to a [`Message`] with a
    /// certain [`Emoji`].
    ///
    /// Refer to [`Channel::reaction_users`] for more information.
    ///
    /// **Note**: Requires the [Read Message History] permission.
    ///
    /// [`Channel::reaction_users`]: enum.Channel.html#method.reaction_users
    /// [`Emoji`]: struct.Emoji.html
    /// [`Message`]: struct.Message.html
    /// [`User`]: struct.User.html
    /// [Read Message History]: permissions/constant.READ_MESSAGE_HISTORY.html
    #[inline]
    pub fn reaction_users<M, R, U>(&self,
                                   message_id: M,
                                   reaction_type: R,
                                   limit: Option<u8>,
                                   after: Option<U>)
                                   -> Result<Vec<User>>
        where M: Into<MessageId>, R: Into<ReactionType>, U: Into<UserId> {
        self.id
            .reaction_users(message_id, reaction_type, limit, after)
    }

    /// Pins a [`Message`] to the channel.
    ///
    /// [`Message`]: struct.Message.html
    #[inline]
    pub fn pin<M: Into<MessageId>>(&self, message_id: M) -> Result<()> { self.id.pin(message_id) }

    /// Retrieves the list of messages that have been pinned in the private
    /// channel.
    #[inline]
    pub fn pins(&self) -> Result<Vec<Message>> { self.id.pins() }

    /// Sends a message with just the given message content in the channel.
    ///
    /// # Errors
    ///
    /// Returns a [`ModelError::MessageTooLong`] if the content of the message
    /// is over the above limit, containing the number of unicode code points
    /// over the limit.
    ///
    /// [`ChannelId`]: ../model/struct.ChannelId.html
    /// [`ModelError::MessageTooLong`]: enum.ModelError.html#variant.MessageTooLong
    #[inline]
    pub fn say<D: ::std::fmt::Display>(&self, content: D) -> Result<Message> { self.id.say(content) }

    /// Sends (a) file(s) along with optional message contents.
    ///
    /// Refer to [`ChannelId::send_files`] for examples and more information.
    ///
    /// The [Attach Files] and [Send Messages] permissions are required.
    ///
    /// **Note**: Message contents must be under 2000 unicode code points.
    ///
    /// # Errors
    ///
    /// If the content of the message is over the above limit, then a
    /// [`ClientError::MessageTooLong`] will be returned, containing the number
    /// of unicode code points over the limit.
    ///
    /// [`ChannelId::send_files`]: struct.ChannelId.html#method.send_files
    /// [`ClientError::MessageTooLong`]: ../client/enum.ClientError.html#variant.MessageTooLong
    /// [Attach Files]: permissions/constant.ATTACH_FILES.html
    /// [Send Messages]: permissions/constant.SEND_MESSAGES.html
    #[inline]
    pub fn send_files<'a, F, T, It: IntoIterator<Item=T>>(&self, files: It, f: F) -> Result<Message>
        where F: FnOnce(CreateMessage) -> CreateMessage, T: Into<AttachmentType<'a>> {
        self.id.send_files(files, f)
    }

    /// Sends a message to the channel with the given content.
    ///
    /// Refer to the documentation for [`CreateMessage`] for more information
    /// regarding message restrictions and requirements.
    ///
    /// # Errors
    ///
    /// Returns a [`ModelError::MessageTooLong`] if the content of the message
    /// is over the above limit, containing the number of unicode code points
    /// over the limit.
    ///
    /// [`ModelError::MessageTooLong`]: enum.ModelError.html#variant.MessageTooLong
    /// [`CreateMessage`]: ../builder/struct.CreateMessage.html
    /// [`Message`]: struct.Message.html
    #[inline]
    pub fn send_message<F: FnOnce(CreateMessage) -> CreateMessage>(&self, f: F) -> Result<Message> {
        self.id.send_message(f)
    }

    /// Unpins a [`Message`] in the channel given by its Id.
    ///
    /// Requires the [Manage Messages] permission.
    ///
    /// [`Message`]: struct.Message.html
    /// [Manage Messages]: permissions/constant.MANAGE_MESSAGES.html
    #[inline]
    pub fn unpin<M: Into<MessageId>>(&self, message_id: M) -> Result<()> {
        self.id.unpin(message_id)
    }
}

impl Display for PrivateChannel {
    /// Formats the private channel, displaying the recipient's username.
    fn fmt(&self, f: &mut Formatter) -> FmtResult {
        f.write_str(&self.recipient.read().name)
    }
}
<commit_msg>Use `AsRef` instead of just `&MessageId`<commit_after>use chrono::{DateTime, FixedOffset};
use model::*;
use std::fmt::{Display, Formatter, Result as FmtResult};
use super::deserialize_single_recipient;

#[cfg(feature = "model")]
use builder::{CreateMessage, GetMessages};
#[cfg(feature = "model")]
use http::AttachmentType;
#[cfg(feature = "model")]
use internal::RwLockExt;

/// A Direct Message text channel with another user.
#[derive(Clone, Debug, Deserialize)]
pub struct PrivateChannel {
    /// The unique Id of the private channel.
    ///
    /// Can be used to calculate the first message's creation date.
    pub id: ChannelId,
    /// The Id of the last message sent.
    pub last_message_id: Option<MessageId>,
    /// Timestamp of the last time a [`Message`] was pinned.
    ///
    /// [`Message`]: struct.Message.html
    pub last_pin_timestamp: Option<DateTime<FixedOffset>>,
    /// Indicator of the type of channel this is.
    ///
    /// This should always be [`ChannelType::Private`].
    ///
    /// [`ChannelType::Private`]: enum.ChannelType.html#variant.Private
    #[serde(rename = "type")]
    pub kind: ChannelType,
    /// The recipient to the private channel.
    #[serde(deserialize_with = "deserialize_single_recipient", rename = "recipients")]
    pub recipient: Arc<RwLock<User>>,
}

#[cfg(feature = "model")]
impl PrivateChannel {
    /// Broadcasts that the current user is typing to the recipient.
    pub fn broadcast_typing(&self) -> Result<()> { self.id.broadcast_typing() }

    /// React to a [`Message`] with a custom [`Emoji`] or unicode character.
    ///
    /// [`Message::react`] may be a more suited method of reacting in most
    /// cases.
    ///
    /// Requires the [Add Reactions] permission, _if_ the current user is the
    /// first user to perform a react with a certain emoji.
    ///
    /// [`Emoji`]: struct.Emoji.html
    /// [`Message`]: struct.Message.html
    /// [`Message::react`]: struct.Message.html#method.react
    /// [Add Reactions]: permissions/constant.ADD_REACTIONS.html
    pub fn create_reaction<M, R>(&self, message_id: M, reaction_type: R) -> Result<()>
        where M: Into<MessageId>, R: Into<ReactionType> {
        self.id.create_reaction(message_id, reaction_type)
    }

    /// Deletes the channel. This does not delete the contents of the channel,
    /// and is equivalent to closing a private channel on the client, which can
    /// be re-opened.
    #[inline]
    pub fn delete(&self) -> Result<Channel> { self.id.delete() }

    /// Deletes all messages by Ids from the given vector in the channel.
    ///
    /// Refer to [`Channel::delete_messages`] for more information.
    ///
    /// Requires the [Manage Messages] permission.
    ///
    /// **Note**: This uses bulk delete endpoint which is not available
    /// for user accounts.
    ///
    /// **Note**: Messages that are older than 2 weeks can't be deleted using
    /// this method.
    ///
    /// [`Channel::delete_messages`]: enum.Channel.html#method.delete_messages
    /// [Manage Messages]: permissions/constant.MANAGE_MESSAGES.html
    #[inline]
    pub fn delete_messages<T: AsRef<MessageId>, It: IntoIterator<Item=T>>(&self, message_ids: It) -> Result<()> {
        self.id.delete_messages(message_ids)
    }

    /// Deletes all permission overrides in the channel from a member
    /// or role.
    ///
    /// **Note**: Requires the [Manage Channel] permission.
    ///
    /// [Manage Channel]: permissions/constant.MANAGE_CHANNELS.html
    #[inline]
    pub fn delete_permission(&self, permission_type: PermissionOverwriteType) -> Result<()> {
        self.id.delete_permission(permission_type)
    }

    /// Deletes the given [`Reaction`] from the channel.
    ///
    /// **Note**: Requires the [Manage Messages] permission, _if_ the current
    /// user did not perform the reaction.
    ///
    /// [`Reaction`]: struct.Reaction.html
    /// [Manage Messages]: permissions/constant.MANAGE_MESSAGES.html
    #[inline]
    pub fn delete_reaction<M, R>(&self,
                                 message_id: M,
                                 user_id: Option<UserId>,
                                 reaction_type: R)
                                 -> Result<()>
        where M: Into<MessageId>, R: Into<ReactionType> {
        self.id.delete_reaction(message_id, user_id, reaction_type)
    }

    /// Edits a [`Message`] in the channel given its Id.
    ///
    /// Message editing preserves all unchanged message data.
    ///
    /// Refer to the documentation for [`CreateMessage`] for more information
    /// regarding message restrictions and requirements.
    ///
    /// **Note**: Requires that the current user be the author of the message.
    ///
    /// # Errors
    ///
    /// Returns a [`ModelError::MessageTooLong`] if the content of the message
    /// is over the [`the limit`], containing the number of unicode code points
    /// over the limit.
    ///
    /// [`ModelError::MessageTooLong`]: enum.ModelError.html#variant.MessageTooLong
    /// [`CreateMessage`]: ../builder/struct.CreateMessage.html
    /// [`Message`]: struct.Message.html
    /// [`the limit`]: ../builder/struct.CreateMessage.html#method.content
    #[inline]
    pub fn edit_message<F, M>(&self, message_id: M, f: F) -> Result<Message>
        where F: FnOnce(CreateMessage) -> CreateMessage, M: Into<MessageId> {
        self.id.edit_message(message_id, f)
    }

    /// Determines if the channel is NSFW.
    ///
    /// Refer to [`utils::is_nsfw`] for more details.
    ///
    /// **Note**: This method is for consistency. This will always return
    /// `false`, due to DMs not being considered NSFW.
    ///
    /// [`utils::is_nsfw`]: ../utils/fn.is_nsfw.html
    #[inline]
    pub fn is_nsfw(&self) -> bool { false }

    /// Gets a message from the channel.
    ///
    /// Requires the [Read Message History] permission.
    ///
    /// [Read Message History]: permissions/constant.READ_MESSAGE_HISTORY.html
    #[inline]
    pub fn message<M: Into<MessageId>>(&self, message_id: M) -> Result<Message> {
        self.id.message(message_id)
    }

    /// Gets messages from the channel.
    ///
    /// Refer to [`Channel::messages`] for more information.
    ///
    /// Requires the [Read Message History] permission.
    ///
    /// [`Channel::messages`]: enum.Channel.html#method.messages
    /// [Read Message History]: permissions/constant.READ_MESSAGE_HISTORY.html
    #[inline]
    pub fn messages<F>(&self, f: F) -> Result<Vec<Message>>
        where F: FnOnce(GetMessages) -> GetMessages {
        self.id.messages(f)
    }

    /// Returns "DM with $username#discriminator".
    pub fn name(&self) -> String { format!("DM with {}", self.recipient.with(|r| r.tag())) }

    /// Gets the list of [`User`]s who have reacted to a [`Message`] with a
    /// certain [`Emoji`].
    ///
    /// Refer to [`Channel::reaction_users`] for more information.
    ///
    /// **Note**: Requires the [Read Message History] permission.
    ///
    /// [`Channel::reaction_users`]: enum.Channel.html#method.reaction_users
    /// [`Emoji`]: struct.Emoji.html
    /// [`Message`]: struct.Message.html
    /// [`User`]: struct.User.html
    /// [Read Message History]: permissions/constant.READ_MESSAGE_HISTORY.html
    #[inline]
    pub fn reaction_users<M, R, U>(&self,
                                   message_id: M,
                                   reaction_type: R,
                                   limit: Option<u8>,
                                   after: Option<U>)
                                   -> Result<Vec<User>>
        where M: Into<MessageId>, R: Into<ReactionType>, U: Into<UserId> {
        self.id
            .reaction_users(message_id, reaction_type, limit, after)
    }

    /// Pins a [`Message`] to the channel.
    ///
    /// [`Message`]: struct.Message.html
    #[inline]
    pub fn pin<M: Into<MessageId>>(&self, message_id: M) -> Result<()> { self.id.pin(message_id) }

    /// Retrieves the list of messages that have been pinned in the private
    /// channel.
    #[inline]
    pub fn pins(&self) -> Result<Vec<Message>> { self.id.pins() }

    /// Sends a message with just the given message content in the channel.
    ///
    /// # Errors
    ///
    /// Returns a [`ModelError::MessageTooLong`] if the content of the message
    /// is over the above limit, containing the number of unicode code points
    /// over the limit.
    ///
    /// [`ChannelId`]: ../model/struct.ChannelId.html
    /// [`ModelError::MessageTooLong`]: enum.ModelError.html#variant.MessageTooLong
    #[inline]
    pub fn say<D: ::std::fmt::Display>(&self, content: D) -> Result<Message> { self.id.say(content) }

    /// Sends (a) file(s) along with optional message contents.
    ///
    /// Refer to [`ChannelId::send_files`] for examples and more information.
    ///
    /// The [Attach Files] and [Send Messages] permissions are required.
    ///
    /// **Note**: Message contents must be under 2000 unicode code points.
    ///
    /// # Errors
    ///
    /// If the content of the message is over the above limit, then a
    /// [`ClientError::MessageTooLong`] will be returned, containing the number
    /// of unicode code points over the limit.
    ///
    /// [`ChannelId::send_files`]: struct.ChannelId.html#method.send_files
    /// [`ClientError::MessageTooLong`]: ../client/enum.ClientError.html#variant.MessageTooLong
    /// [Attach Files]: permissions/constant.ATTACH_FILES.html
    /// [Send Messages]: permissions/constant.SEND_MESSAGES.html
    #[inline]
    pub fn send_files<'a, F, T, It: IntoIterator<Item=T>>(&self, files: It, f: F) -> Result<Message>
        where F: FnOnce(CreateMessage) -> CreateMessage, T: Into<AttachmentType<'a>> {
        self.id.send_files(files, f)
    }

    /// Sends a message to the channel with the given content.
    ///
    /// Refer to the documentation for [`CreateMessage`] for more information
    /// regarding message restrictions and requirements.
    ///
    /// # Errors
    ///
    /// Returns a [`ModelError::MessageTooLong`] if the content of the message
    /// is over the above limit, containing the number of unicode code points
    /// over the limit.
    ///
    /// [`ModelError::MessageTooLong`]: enum.ModelError.html#variant.MessageTooLong
    /// [`CreateMessage`]: ../builder/struct.CreateMessage.html
    /// [`Message`]: struct.Message.html
    #[inline]
    pub fn send_message<F: FnOnce(CreateMessage) -> CreateMessage>(&self, f: F) -> Result<Message> {
        self.id.send_message(f)
    }

    /// Unpins a [`Message`] in the channel given by its Id.
    ///
    /// Requires the [Manage Messages] permission.
    ///
    /// [`Message`]: struct.Message.html
    /// [Manage Messages]: permissions/constant.MANAGE_MESSAGES.html
    #[inline]
    pub fn unpin<M: Into<MessageId>>(&self, message_id: M) -> Result<()> {
        self.id.unpin(message_id)
    }
}

impl Display for PrivateChannel {
    /// Formats the private channel, displaying the recipient's username.
    fn fmt(&self, f: &mut Formatter) -> FmtResult {
        f.write_str(&self.recipient.read().name)
    }
}
<|endoftext|>


<commit_before>use autograd::{Function, FuncIntf, FuncDelegate, Variable, VarKindList, FIWrap};
use macros::*;
use tensor::{TensorKindList, OptTensorKindList};
use ::*;

pub struct Threshold {
    delegate: Function,
    threshold: f32,
    value: f32,
    inplace: bool,
}

impl Threshold {
    pub fn new(threshold: f32, value: f32, inplace: bool) -> FIWrap<Self> {
        FIWrap::new(Threshold {
                        delegate: Function::new(),
                        threshold: threshold,
                        value: value,
                        inplace: inplace,
                    })
    }
}

impl_func_delegate!(Threshold);

impl FuncIntf for Threshold {
    fn forward(&mut self, input: &mut TensorKindList) -> TensorKindList {
        unimplemented!()
    }
    fn backward(&mut self, input: &mut OptTensorKindList) -> OptTensorKindList {
        unimplemented!()
    }
}
<commit_msg>implement forward threshold<commit_after>use autograd::{Function, FuncIntf, FuncDelegate, FIWrap};
use tensor::{TensorKindList, OptTensorKindList, New};

pub struct Threshold {
    delegate: Function,
    threshold: f32,
    value: f32,
    inplace: bool,
}

impl Threshold {
    pub fn new(threshold: f32, value: f32, inplace: bool) -> FIWrap<Self> {
        if inplace {
            panic!("in-place processing requires value ({}) to not \
                    exceed threshold ({})",
                   value,
                   threshold);
        }

        FIWrap::new(Threshold {
                        delegate: Function::new(),
                        threshold: threshold,
                        value: value,
                        inplace: inplace,
                    })
    }
}

impl_func_delegate!(Threshold);

impl FuncIntf for Threshold {
    fn forward(&mut self, input_: &mut TensorKindList) -> TensorKindList {
        let backend = input_[0].backend();
        let mut output = if self.inplace {
            self.mark_dirty(input_);
            input_[0].clone()
        } else {
            input_[0].new_(())
        };
        self.save_for_backward(input_);
        backend.Threshold_updateOutput(&mut input_[0],
                                       &mut output,
                                       self.threshold,
                                       self.value,
                                       self.inplace);
        vec![output]
    }
    fn backward(&mut self, input: &mut OptTensorKindList) -> OptTensorKindList {
        unimplemented!()
    }
}
<|endoftext|>


<commit_before>// Copyright 2015-2016 ZFilexfer Developers. See the COPYRIGHT file at the
// top-level directory of this distribution and at
// https://intecture.io/COPYRIGHT.
//
// Licensed under the Mozilla Public License 2.0 <LICENSE or
// https://www.tldrlegal.com/l/mpl-2.0>. This file may not be copied,
// modified, or distributed except according to those terms.

use chunk::Chunk;
use czmq::{ZMsg, ZSock, ZSys};
use error::{Error, Result};
use std::rc::Rc;
use std::sync::{Arc, RwLock};
use std::thread::{JoinHandle, spawn};
use std::time::Instant;

#[cfg(not(test))]
const CHUNK_TIMEOUT: u64 = 60;
#[cfg(test)]
const CHUNK_TIMEOUT: u64 = 1;

pub struct Arbitrator {
    router: Rc<ZSock>,
    queue: Arc<RwLock<Vec<TimedChunk>>>,
    timer_handle: Option<JoinHandle<()>>,
    timer_comm: ZSock,
    slots: u32,
}

impl Drop for Arbitrator {
    fn drop(&mut self) {
        self.timer_comm.signal(0).unwrap();
        if self.timer_handle.is_some() {
            self.timer_handle.take().unwrap().join().unwrap();
        }
    }
}

impl Arbitrator {
    pub fn new(router: Rc<ZSock>, upload_slots: u32) -> Result<Arbitrator> {
        let (comm_front, comm_back) = try!(ZSys::create_pipe());
        comm_front.set_sndtimeo(Some(1000));
        comm_back.set_rcvtimeo(Some(1000)); // Remember that this timeout controls the Timer loop speed!

        let lock = Arc::new(RwLock::new(Vec::new()));
        let timer = try!(Timer::new(comm_back, lock.clone()));

        Ok(Arbitrator {
            router: router,
            queue: lock,
            timer_handle: Some(spawn(move|| timer.run())),
            timer_comm: comm_front,
            slots: upload_slots,
        })
    }

    pub fn queue(&mut self, chunk: &Chunk, router_id: &[u8]) -> Result<()> {
        let timed_chunk = TimedChunk::new(router_id, chunk.get_index());
        {
            let mut writer = self.queue.write().unwrap();
            writer.push(timed_chunk);
        }

        try!(self.request());
        Ok(())
    }

    pub fn release(&mut self, chunk: &Chunk, router_id: &[u8]) -> Result<()> {
        let router_id = router_id.to_vec();
        {
            let mut queue = self.queue.write().unwrap();
            let mut index: Option<usize> = None;
            let mut x = 0;
            for c in queue.iter_mut() {
                if c.router_id == router_id && c.index == chunk.get_index() {
                    index = Some(x);
                    break;
                }

                x += 1;
            }

            match index {
                Some(i) => {
                    queue.remove(i);
                    self.slots += 1;
                },
                None => return Err(Error::ChunkIndex),
            }
        }

        try!(self.request());
        Ok(())
    }

    fn request(&mut self) -> Result<()> {
        for chunk in self.queue.write().unwrap().iter_mut() {
            if self.slots == 0 {
                break;
            }

            if !chunk.is_started() {
                self.slots -= 1;

                let msg = ZMsg::new();
                try!(msg.addbytes(&chunk.router_id));
                try!(msg.addstr("CHUNK"));
                try!(msg.addstr(&chunk.index.to_string()));
                try!(msg.send(&*self.router));

                chunk.start();
            }
        }

        Ok(())
    }
}

struct Timer {
    chunks: Arc<RwLock<Vec<TimedChunk>>>,
    sink: ZSock,
    comm: ZSock,
}

impl Timer {
    fn new(comm: ZSock, chunks: Arc<RwLock<Vec<TimedChunk>>>) -> Result<Timer> {
        Ok(Timer {
            chunks: chunks,
            sink: try!(ZSock::new_push(">inproc://zfilexfer_sink")),
            comm: comm,
        })
    }

    fn run(self) {
        loop {
            // Terminate on signal
            if self.comm.wait().is_ok() {
                break;
            }

            for chunk in self.chunks.read().unwrap().iter() {
                if chunk.is_expired() {
                    let msg = ZMsg::new();
                    msg.addbytes(&chunk.router_id).unwrap();
                    msg.addstr(&chunk.index.to_string()).unwrap();
                    msg.addstr("0").unwrap();
                    msg.send(&self.sink).unwrap();
                }
            }
        }
    }
}

struct TimedChunk {
    router_id: Vec<u8>,
    index: u64,
    timestamp: Option<Instant>,
}

impl TimedChunk {
    fn new(router_id: &[u8], index: u64) -> TimedChunk {
        TimedChunk {
            router_id: router_id.to_vec(),
            index: index,
            timestamp: None,
        }
    }

    fn start(&mut self) {
        self.timestamp = Some(Instant::now());
    }

    fn is_started(&self) -> bool {
        self.timestamp.is_some()
    }

    fn is_expired(&self) -> bool {
        if self.timestamp.is_some() {
            self.timestamp.as_ref().unwrap().elapsed().as_secs() >= CHUNK_TIMEOUT
        } else {
            false
        }
    }
}

#[cfg(test)]
mod tests {
    use chunk::Chunk;
    use czmq::{ZMsg, ZSock, ZSockType, ZSys};
    use std::rc::Rc;
    use std::sync::{Arc, RwLock};
    use std::thread::{sleep, spawn};
    use std::time::{Duration, Instant};
    use super::*;
    use super::{TimedChunk, Timer};

    #[test]
    fn test_arbitrator_new() {
        ZSys::init();

        assert!(Arbitrator::new(Rc::new(ZSock::new(ZSockType::ROUTER)), 0).is_ok());
    }

    #[test]
    fn test_arbitrator_queue_release() {
        ZSys::init();

        let chunk = Chunk::test_new("/path/to/file", 0);

        let mut arbitrator = Arbitrator::new(Rc::new(ZSock::new(ZSockType::ROUTER)), 1).unwrap();
        assert!(arbitrator.queue(&chunk, "abc".as_bytes()).is_ok());
        assert_eq!(arbitrator.queue.read().unwrap().len(), 1);
        assert_eq!(arbitrator.slots, 0);
        assert!(arbitrator.release(&chunk, "abc".as_bytes()).is_ok());
        assert_eq!(arbitrator.queue.read().unwrap().len(), 0);
        assert_eq!(arbitrator.slots, 1);
    }

    #[test]
    fn test_arbitrator_request() {
        ZSys::init();

        let (client, router) = ZSys::create_pipe().unwrap();
        client.set_rcvtimeo(Some(500));

        let (comm, thread) = ZSys::create_pipe().unwrap();

        let chunks = vec![
            TimedChunk::new("abc".as_bytes(), 0),
            TimedChunk::new("abc".as_bytes(), 1),
            TimedChunk::new("abc".as_bytes(), 2),
            TimedChunk::new("def".as_bytes(), 0),
            TimedChunk::new("def".as_bytes(), 1),
            TimedChunk::new("def".as_bytes(), 2),
        ];

        {
            let mut arbitrator = Arbitrator {
                router: Rc::new(router),
                queue: Arc::new(RwLock::new(chunks)),
                timer_handle: None,
                timer_comm: comm,
                slots: 3,
            };

            arbitrator.request().unwrap();

            for x in 0..3 {
                let msg = ZMsg::recv(&client).unwrap();
                assert_eq!(&msg.popstr().unwrap().unwrap(), "abc");
                assert_eq!(&msg.popstr().unwrap().unwrap(), "CHUNK");
                assert_eq!(msg.popstr().unwrap().unwrap(), x.to_string());
            }

            assert!(client.recv_str().is_err());

            let chunk = Chunk::test_new("/path/to/file", 0);
            arbitrator.release(&chunk, "abc".as_bytes()).unwrap();
            arbitrator.request().unwrap();

            let msg = ZMsg::recv(&client).unwrap();
            assert_eq!(&msg.popstr().unwrap().unwrap(), "def");
            assert_eq!(&msg.popstr().unwrap().unwrap(), "CHUNK");
            assert_eq!(msg.popstr().unwrap().unwrap(), "0");
        }

        thread.wait().unwrap();
    }

    #[test]
    fn test_timer_new() {
        ZSys::init();

        assert!(Timer::new(ZSock::new(ZSockType::REQ), Arc::new(RwLock::new(Vec::new()))).is_ok());
    }

    #[test]
    fn test_timer_run() {
        ZSys::init();

        let (client, server) = ZSys::create_pipe().unwrap();
        let (comm, thread) = ZSys::create_pipe().unwrap();
        client.set_rcvtimeo(Some(1500));
        thread.set_rcvtimeo(Some(1000));

        let mut c = TimedChunk::new("abc".as_bytes(), 0);
        c.start();

        let timer = Timer {
            chunks: Arc::new(RwLock::new(vec![
                c,
            ])),
            sink: server,
            comm: thread,
        };
        let handle = spawn(|| timer.run());

        let msg = ZMsg::recv(&client).unwrap();
        assert_eq!(msg.popstr().unwrap().unwrap(), "abc");
        assert_eq!(msg.popstr().unwrap().unwrap(), "0");
        assert_eq!(msg.popstr().unwrap().unwrap(), "0");

        comm.signal(0).unwrap();
        handle.join().unwrap();
    }

    #[test]
    fn test_chunk_is_expired() {
        let timed = TimedChunk {
            router_id: vec![97, 98, 99],
            index: 0,
            timestamp: Some(Instant::now()),
        };

        sleep(Duration::from_secs(1));

        assert!(timed.is_expired());
    }
}
<commit_msg>Handle system signals in child thread<commit_after>// Copyright 2015-2016 ZFilexfer Developers. See the COPYRIGHT file at the
// top-level directory of this distribution and at
// https://intecture.io/COPYRIGHT.
//
// Licensed under the Mozilla Public License 2.0 <LICENSE or
// https://www.tldrlegal.com/l/mpl-2.0>. This file may not be copied,
// modified, or distributed except according to those terms.

use chunk::Chunk;
use czmq::{ZMsg, ZSock, ZSys};
use error::{Error, Result};
use std::rc::Rc;
use std::sync::{Arc, RwLock};
use std::thread::{JoinHandle, spawn};
use std::time::Instant;

#[cfg(not(test))]
const CHUNK_TIMEOUT: u64 = 60;
#[cfg(test)]
const CHUNK_TIMEOUT: u64 = 1;

pub struct Arbitrator {
    router: Rc<ZSock>,
    queue: Arc<RwLock<Vec<TimedChunk>>>,
    timer_handle: Option<JoinHandle<()>>,
    timer_comm: ZSock,
    slots: u32,
}

impl Drop for Arbitrator {
    fn drop(&mut self) {
        // Ignore failure as it means the thread has already
        // terminated.
        let _ = self.timer_comm.signal(0);
        if self.timer_handle.is_some() {
            self.timer_handle.take().unwrap().join().unwrap();
        }
    }
}

impl Arbitrator {
    pub fn new(router: Rc<ZSock>, upload_slots: u32) -> Result<Arbitrator> {
        let (comm_front, comm_back) = try!(ZSys::create_pipe());
        comm_front.set_sndtimeo(Some(1000));
        comm_back.set_rcvtimeo(Some(1000)); // Remember that this timeout controls the Timer loop speed!

        let lock = Arc::new(RwLock::new(Vec::new()));
        let timer = try!(Timer::new(comm_back, lock.clone()));

        Ok(Arbitrator {
            router: router,
            queue: lock,
            timer_handle: Some(spawn(move|| timer.run())),
            timer_comm: comm_front,
            slots: upload_slots,
        })
    }

    pub fn queue(&mut self, chunk: &Chunk, router_id: &[u8]) -> Result<()> {
        let timed_chunk = TimedChunk::new(router_id, chunk.get_index());
        {
            let mut writer = self.queue.write().unwrap();
            writer.push(timed_chunk);
        }

        try!(self.request());
        Ok(())
    }

    pub fn release(&mut self, chunk: &Chunk, router_id: &[u8]) -> Result<()> {
        let router_id = router_id.to_vec();
        {
            let mut queue = self.queue.write().unwrap();
            let mut index: Option<usize> = None;
            let mut x = 0;
            for c in queue.iter_mut() {
                if c.router_id == router_id && c.index == chunk.get_index() {
                    index = Some(x);
                    break;
                }

                x += 1;
            }

            match index {
                Some(i) => {
                    queue.remove(i);
                    self.slots += 1;
                },
                None => return Err(Error::ChunkIndex),
            }
        }

        try!(self.request());
        Ok(())
    }

    fn request(&mut self) -> Result<()> {
        for chunk in self.queue.write().unwrap().iter_mut() {
            if self.slots == 0 {
                break;
            }

            if !chunk.is_started() {
                self.slots -= 1;

                let msg = ZMsg::new();
                try!(msg.addbytes(&chunk.router_id));
                try!(msg.addstr("CHUNK"));
                try!(msg.addstr(&chunk.index.to_string()));
                try!(msg.send(&*self.router));

                chunk.start();
            }
        }

        Ok(())
    }
}

struct Timer {
    chunks: Arc<RwLock<Vec<TimedChunk>>>,
    sink: ZSock,
    comm: ZSock,
}

impl Timer {
    fn new(comm: ZSock, chunks: Arc<RwLock<Vec<TimedChunk>>>) -> Result<Timer> {
        Ok(Timer {
            chunks: chunks,
            sink: try!(ZSock::new_push(">inproc://zfilexfer_sink")),
            comm: comm,
        })
    }

    fn run(self) {
        loop {
            // Terminate on ZSock signal or system signal (SIGTERM)
            if self.comm.wait().is_ok() || ZSys::is_interrupted() {
                break;
            }

            for chunk in self.chunks.read().unwrap().iter() {
                if chunk.is_expired() {
                    let msg = ZMsg::new();
                    msg.addbytes(&chunk.router_id).unwrap();
                    msg.addstr(&chunk.index.to_string()).unwrap();
                    msg.addstr("0").unwrap();
                    msg.send(&self.sink).unwrap();
                }
            }
        }
    }
}

struct TimedChunk {
    router_id: Vec<u8>,
    index: u64,
    timestamp: Option<Instant>,
}

impl TimedChunk {
    fn new(router_id: &[u8], index: u64) -> TimedChunk {
        TimedChunk {
            router_id: router_id.to_vec(),
            index: index,
            timestamp: None,
        }
    }

    fn start(&mut self) {
        self.timestamp = Some(Instant::now());
    }

    fn is_started(&self) -> bool {
        self.timestamp.is_some()
    }

    fn is_expired(&self) -> bool {
        if self.timestamp.is_some() {
            self.timestamp.as_ref().unwrap().elapsed().as_secs() >= CHUNK_TIMEOUT
        } else {
            false
        }
    }
}

#[cfg(test)]
mod tests {
    use chunk::Chunk;
    use czmq::{ZMsg, ZSock, ZSockType, ZSys};
    use std::rc::Rc;
    use std::sync::{Arc, RwLock};
    use std::thread::{sleep, spawn};
    use std::time::{Duration, Instant};
    use super::*;
    use super::{TimedChunk, Timer};

    #[test]
    fn test_arbitrator_new() {
        ZSys::init();

        assert!(Arbitrator::new(Rc::new(ZSock::new(ZSockType::ROUTER)), 0).is_ok());
    }

    #[test]
    fn test_arbitrator_queue_release() {
        ZSys::init();

        let chunk = Chunk::test_new("/path/to/file", 0);

        let mut arbitrator = Arbitrator::new(Rc::new(ZSock::new(ZSockType::ROUTER)), 1).unwrap();
        assert!(arbitrator.queue(&chunk, "abc".as_bytes()).is_ok());
        assert_eq!(arbitrator.queue.read().unwrap().len(), 1);
        assert_eq!(arbitrator.slots, 0);
        assert!(arbitrator.release(&chunk, "abc".as_bytes()).is_ok());
        assert_eq!(arbitrator.queue.read().unwrap().len(), 0);
        assert_eq!(arbitrator.slots, 1);
    }

    #[test]
    fn test_arbitrator_request() {
        ZSys::init();

        let (client, router) = ZSys::create_pipe().unwrap();
        client.set_rcvtimeo(Some(500));

        let (comm, thread) = ZSys::create_pipe().unwrap();

        let chunks = vec![
            TimedChunk::new("abc".as_bytes(), 0),
            TimedChunk::new("abc".as_bytes(), 1),
            TimedChunk::new("abc".as_bytes(), 2),
            TimedChunk::new("def".as_bytes(), 0),
            TimedChunk::new("def".as_bytes(), 1),
            TimedChunk::new("def".as_bytes(), 2),
        ];

        {
            let mut arbitrator = Arbitrator {
                router: Rc::new(router),
                queue: Arc::new(RwLock::new(chunks)),
                timer_handle: None,
                timer_comm: comm,
                slots: 3,
            };

            arbitrator.request().unwrap();

            for x in 0..3 {
                let msg = ZMsg::recv(&client).unwrap();
                assert_eq!(&msg.popstr().unwrap().unwrap(), "abc");
                assert_eq!(&msg.popstr().unwrap().unwrap(), "CHUNK");
                assert_eq!(msg.popstr().unwrap().unwrap(), x.to_string());
            }

            assert!(client.recv_str().is_err());

            let chunk = Chunk::test_new("/path/to/file", 0);
            arbitrator.release(&chunk, "abc".as_bytes()).unwrap();
            arbitrator.request().unwrap();

            let msg = ZMsg::recv(&client).unwrap();
            assert_eq!(&msg.popstr().unwrap().unwrap(), "def");
            assert_eq!(&msg.popstr().unwrap().unwrap(), "CHUNK");
            assert_eq!(msg.popstr().unwrap().unwrap(), "0");
        }

        thread.wait().unwrap();
    }

    #[test]
    fn test_timer_new() {
        ZSys::init();

        assert!(Timer::new(ZSock::new(ZSockType::REQ), Arc::new(RwLock::new(Vec::new()))).is_ok());
    }

    #[test]
    fn test_timer_run() {
        ZSys::init();

        let (client, server) = ZSys::create_pipe().unwrap();
        let (comm, thread) = ZSys::create_pipe().unwrap();
        client.set_rcvtimeo(Some(1500));
        thread.set_rcvtimeo(Some(1000));

        let mut c = TimedChunk::new("abc".as_bytes(), 0);
        c.start();

        let timer = Timer {
            chunks: Arc::new(RwLock::new(vec![
                c,
            ])),
            sink: server,
            comm: thread,
        };
        let handle = spawn(|| timer.run());

        let msg = ZMsg::recv(&client).unwrap();
        assert_eq!(msg.popstr().unwrap().unwrap(), "abc");
        assert_eq!(msg.popstr().unwrap().unwrap(), "0");
        assert_eq!(msg.popstr().unwrap().unwrap(), "0");

        comm.signal(0).unwrap();
        handle.join().unwrap();
    }

    #[test]
    fn test_chunk_is_expired() {
        let timed = TimedChunk {
            router_id: vec![97, 98, 99],
            index: 0,
            timestamp: Some(Instant::now()),
        };

        sleep(Duration::from_secs(1));

        assert!(timed.is_expired());
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

use std::collections::{vec_deque, VecDeque};
use std::error::Error;
use std::str;

use bam;
use bam::Read;

/// A buffer for BAM records. This allows access regions in a sorted BAM file while iterating
/// over it in a single pass.
/// The buffer is implemented as a ringbuffer, such that extension or movement to the right has
/// linear complexity. The buffer makes use of indexed random access. Hence, when fetching a
/// region at the very end of the BAM, everything before is omitted without cost.
#[derive(Debug)]
pub struct RecordBuffer {
    reader: bam::IndexedReader,
    inner: VecDeque<bam::Record>,
    overflow: Option<bam::Record>
}

unsafe impl Sync for RecordBuffer {}
unsafe impl Send for RecordBuffer {}

impl RecordBuffer {
    /// Create a new `RecordBuffer`.
    pub fn new(bam: bam::IndexedReader) -> Self {
        RecordBuffer {
            reader: bam,
            inner: VecDeque::new(),
            overflow: None
        }
    }

    /// Return end position of buffer.
    fn end(&self) -> Option<u32> {
        self.inner.back().map(|rec| rec.pos() as u32)
    }

    fn tid(&self) -> Option<i32> {
        self.inner.back().map(|rec| rec.tid())
    }

    /// Fill buffer at the given interval. The start coordinate has to be left of
    /// the start coordinate of any previous `fill` operation.
    /// Coordinates are 0-based, and end is exclusive.
    /// Returns tuple with numbers of added and deleted records since the previous fetch.
    #[allow(unused_assignments)]  // TODO this is needed because rustc thinks that deleted is unused
    pub fn fetch(&mut self, chrom: &[u8], start: u32, end: u32) -> Result<(usize, usize), Box<Error>> {
        let mut added = 0;
        // move overflow from last fetch into ringbuffer
        if self.overflow.is_some() {
            added += 1;
            self.inner.push_back(self.overflow.take().unwrap());
        }

        if let Some(tid) = self.reader.header.tid(chrom) {
            let mut deleted = 0;
            let window_start = start;
            if self.inner.is_empty() || self.end().unwrap() < window_start || self.tid().unwrap() != tid as i32 {
                let end = self.reader.header.target_len(tid).unwrap();
                self.reader.fetch(tid, window_start, end)?;
                deleted = self.inner.len();
                self.inner.clear();
            } else {
                // remove records too far left
                let to_remove = self.inner.iter().take_while(|rec| rec.pos() < window_start as i32).count();
                for _ in 0..to_remove {
                    self.inner.pop_front();
                }
                deleted = to_remove;
            }

            // extend to the right
            loop {
                let mut record = bam::Record::new();
                if let Err(e) = self.reader.read(&mut record) {
                    if e.is_eof() {
                        break;
                    }
                    return Err(Box::new(e));
                }

                if record.is_unmapped() {
                    continue;
                }

                let pos = record.pos();
                if pos >= end as i32 {
                    self.overflow = Some(record);
                    break;
                } else {
                    self.inner.push_back(record);
                    added += 1;
                }
            }

            Ok((added, deleted))
        } else {
            Err(Box::new(RecordBufferError::UnknownSequence(str::from_utf8(chrom).unwrap().to_owned())))
        }
    }

    /// Iterate over records that have been fetched with `fetch`.
    pub fn iter(&self) -> vec_deque::Iter<bam::Record> {
        self.inner.iter()
    }

    /// Iterate over mutable references to records that have been fetched with `fetch`.
    pub fn iter_mut(&mut self) -> vec_deque::IterMut<bam::Record> {
        self.inner.iter_mut()
    }

    pub fn len(&self) -> usize {
        self.inner.len()
    }
}

quick_error! {
    #[derive(Debug, Clone)]
    pub enum RecordBufferError {
        UnknownSequence(chrom: String) {
            description("unknown sequence")
            display("sequence {} cannot be found in BAM", chrom)
        }
    }
}

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

    #[test]
    fn test_buffer() {
        let reader = bam::IndexedReader::from_path(&"test/test.bam").unwrap();

        let mut buffer = RecordBuffer::new(reader);

        buffer.fetch(b"CHROMOSOME_I", 1, 5).unwrap();
        {
            let records = buffer.iter().collect_vec();
            assert_eq!(records.len(), 6);
            assert_eq!(records[0].pos(), 1);
            assert_eq!(records[1].pos(), 1);
            assert_eq!(records[2].pos(), 1);
            assert_eq!(records[3].pos(), 1);
            assert_eq!(records[4].pos(), 1);
            assert_eq!(records[5].pos(), 1);
        }
    }
}
<commit_msg>Allow going left in bam buffer.<commit_after>// Copyright 2017 Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

use std::collections::{vec_deque, VecDeque};
use std::error::Error;
use std::str;

use bam;
use bam::Read;

/// A buffer for BAM records. This allows access regions in a sorted BAM file while iterating
/// over it in a single pass.
/// The buffer is implemented as a ringbuffer, such that extension or movement to the right has
/// linear complexity. The buffer makes use of indexed random access. Hence, when fetching a
/// region at the very end of the BAM, everything before is omitted without cost.
#[derive(Debug)]
pub struct RecordBuffer {
    reader: bam::IndexedReader,
    inner: VecDeque<bam::Record>,
    overflow: Option<bam::Record>
}

unsafe impl Sync for RecordBuffer {}
unsafe impl Send for RecordBuffer {}

impl RecordBuffer {
    /// Create a new `RecordBuffer`.
    pub fn new(bam: bam::IndexedReader) -> Self {
        RecordBuffer {
            reader: bam,
            inner: VecDeque::new(),
            overflow: None
        }
    }

    /// Return start position of buffer
    fn start(&self) -> Option<u32> {
        self.inner.front().map(|rec| rec.pos() as u32)
    }

    /// Return end position of buffer.
    fn end(&self) -> Option<u32> {
        self.inner.back().map(|rec| rec.pos() as u32)
    }

    fn tid(&self) -> Option<i32> {
        self.inner.back().map(|rec| rec.tid())
    }

    /// Fill buffer at the given interval. If the start coordinate is left of
    /// the previous start coordinate, this will use an additional BAM fetch IO operation.
    /// Coordinates are 0-based, and end is exclusive.
    /// Returns tuple with numbers of added and deleted records since the previous fetch.
    #[allow(unused_assignments)]  // TODO this is needed because rustc thinks that deleted is unused
    pub fn fetch(&mut self, chrom: &[u8], start: u32, end: u32) -> Result<(usize, usize), Box<Error>> {
        let mut added = 0;
        // move overflow from last fetch into ringbuffer
        if self.overflow.is_some() {
            added += 1;
            self.inner.push_back(self.overflow.take().unwrap());
        }

        if let Some(tid) = self.reader.header.tid(chrom) {
            let mut deleted = 0;
            let window_start = start;
            if self.inner.is_empty() ||
               self.end().unwrap() < window_start ||
               self.tid().unwrap() != tid as i32 ||
               self.start().unwrap() > window_start
            {
                let end = self.reader.header.target_len(tid).unwrap();
                self.reader.fetch(tid, window_start, end)?;
                deleted = self.inner.len();
                self.inner.clear();
            } else {
                // remove records too far left
                let to_remove = self.inner.iter().take_while(|rec| rec.pos() < window_start as i32).count();
                for _ in 0..to_remove {
                    self.inner.pop_front();
                }
                deleted = to_remove;
            }

            // extend to the right
            loop {
                let mut record = bam::Record::new();
                if let Err(e) = self.reader.read(&mut record) {
                    if e.is_eof() {
                        break;
                    }
                    return Err(Box::new(e));
                }

                if record.is_unmapped() {
                    continue;
                }

                let pos = record.pos();
                if pos >= end as i32 {
                    self.overflow = Some(record);
                    break;
                } else {
                    self.inner.push_back(record);
                    added += 1;
                }
            }

            Ok((added, deleted))
        } else {
            Err(Box::new(RecordBufferError::UnknownSequence(str::from_utf8(chrom).unwrap().to_owned())))
        }
    }

    /// Iterate over records that have been fetched with `fetch`.
    pub fn iter(&self) -> vec_deque::Iter<bam::Record> {
        self.inner.iter()
    }

    /// Iterate over mutable references to records that have been fetched with `fetch`.
    pub fn iter_mut(&mut self) -> vec_deque::IterMut<bam::Record> {
        self.inner.iter_mut()
    }

    pub fn len(&self) -> usize {
        self.inner.len()
    }
}

quick_error! {
    #[derive(Debug, Clone)]
    pub enum RecordBufferError {
        UnknownSequence(chrom: String) {
            description("unknown sequence")
            display("sequence {} cannot be found in BAM", chrom)
        }
    }
}

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

    #[test]
    fn test_buffer() {
        let reader = bam::IndexedReader::from_path(&"test/test.bam").unwrap();

        let mut buffer = RecordBuffer::new(reader);

        buffer.fetch(b"CHROMOSOME_I", 1, 5).unwrap();
        {
            let records = buffer.iter().collect_vec();
            assert_eq!(records.len(), 6);
            assert_eq!(records[0].pos(), 1);
            assert_eq!(records[1].pos(), 1);
            assert_eq!(records[2].pos(), 1);
            assert_eq!(records[3].pos(), 1);
            assert_eq!(records[4].pos(), 1);
            assert_eq!(records[5].pos(), 1);
        }
    }
}
<|endoftext|>


<commit_before>#![warn(clippy::all, rust_2018_idioms)]
#![allow(clippy::unknown_clippy_lints)]

use cargo_registry::{boot, App, Env};
use std::{
    borrow::Cow,
    fs::File,
    sync::{mpsc::channel, Arc, Mutex},
    thread,
    time::Duration,
};

use civet::Server as CivetServer;
use conduit_hyper::Service;
use futures_util::future::FutureExt;
use reqwest::blocking::Client;
use sentry::{ClientOptions, IntoDsn};

const CORE_THREADS: usize = 4;

#[allow(clippy::large_enum_variant)]
enum Server {
    Civet(CivetServer),
    Hyper(tokio::runtime::Runtime, tokio::task::JoinHandle<()>),
}

use Server::*;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let _sentry = dotenv::var("SENTRY_DSN_API")
        .ok()
        .into_dsn()
        .expect("SENTRY_DSN_API is not a valid Sentry DSN value")
        .map(|dsn| {
            let mut opts = ClientOptions::from(dsn);
            opts.environment = Some(
                dotenv::var("SENTRY_ENV_API")
                    .map(Cow::Owned)
                    .expect("SENTRY_ENV_API must be set when using SENTRY_DSN_API"),
            );
            opts.release = option_env!("CRATES_IO_GIT_COMMIT").map(Into::into);

            sentry::init(opts)
        });

    // Initialize logging
    env_logger::init();

    let config = cargo_registry::Config::default();
    let client = Client::new();

    let app = App::new(config.clone(), Some(client));
    let app = cargo_registry::build_handler(Arc::new(app));

    // On every server restart, ensure the categories available in the database match
    // the information in *src/categories.toml*.
    let categories_toml = include_str!("../boot/categories.toml");
    boot::categories::sync(categories_toml).unwrap();

    let heroku = dotenv::var("HEROKU").is_ok();
    let fastboot = dotenv::var("USE_FASTBOOT").is_ok();

    let port = if heroku {
        8888
    } else {
        dotenv::var("PORT")
            .ok()
            .and_then(|s| s.parse().ok())
            .unwrap_or(8888)
    };
    let threads = dotenv::var("SERVER_THREADS")
        .map(|s| s.parse().expect("SERVER_THREADS was not a valid number"))
        .unwrap_or_else(|_| {
            if config.env == Env::Development {
                5
            } else {
                // A large default because this can be easily changed via env and in production we
                // want the logging middleware to accurately record the start time.
                500
            }
        });

    let server = if dotenv::var("WEB_USE_CIVET").is_err() {
        use tokio::io::AsyncWriteExt;
        use tokio::signal::unix::{signal, SignalKind};

        println!("Booting with a hyper based server");

        let mut rt = tokio::runtime::Builder::new()
            .threaded_scheduler()
            .enable_all()
            .core_threads(CORE_THREADS)
            .max_threads(threads as usize + CORE_THREADS)
            .build()
            .unwrap();

        let handler = Arc::new(conduit_hyper::BlockingHandler::new(app));
        let make_service =
            hyper::service::make_service_fn(move |socket: &hyper::server::conn::AddrStream| {
                let addr = socket.remote_addr();
                let handler = handler.clone();
                async move { Service::from_blocking(handler, addr) }
            });

        let addr = ([127, 0, 0, 1], port).into();
        #[allow(clippy::async_yields_async)]
        let server = rt.block_on(async { hyper::Server::bind(&addr).serve(make_service) });

        let mut sig_int = rt.block_on(async { signal(SignalKind::interrupt()) })?;
        let mut sig_term = rt.block_on(async { signal(SignalKind::terminate()) })?;

        let server = server.with_graceful_shutdown(async move {
            // Wait for either signal
            futures_util::select! {
                _ = sig_int.recv().fuse() => (),
                _ = sig_term.recv().fuse() => (),
            };
            let mut stdout = tokio::io::stdout();
            stdout.write_all(b"Starting graceful shutdown\n").await.ok();
        });

        let server = rt.spawn(async { server.await.unwrap() });
        Hyper(rt, server)
    } else {
        println!("Booting with a civet based server");
        let mut cfg = civet::Config::new();
        cfg.port(port).threads(threads).keep_alive(true);
        Civet(CivetServer::start(cfg, app).unwrap())
    };

    println!("listening on port {}", port);

    // Give tokio a chance to spawn the first worker thread
    thread::sleep(Duration::from_millis(10));

    // Creating this file tells heroku to tell nginx that the application is ready
    // to receive traffic.
    if heroku {
        let path = if fastboot {
            "/tmp/backend-initialized"
        } else {
            "/tmp/app-initialized"
        };
        println!("Writing to {}", path);
        File::create(path).unwrap();
    }

    // Block the main thread until the server has shutdown
    match server {
        Hyper(mut rt, server) => {
            rt.block_on(async { server.await.unwrap() });
        }
        Civet(server) => {
            let (tx, rx) = channel::<()>();
            ctrlc_handler(move || tx.send(()).unwrap_or(()));
            rx.recv().unwrap();
            drop(server);
        }
    }

    println!("Server has gracefully shutdown!");
    Ok(())
}

fn ctrlc_handler<F>(f: F)
where
    F: FnOnce() + Send + 'static,
{
    let call_once = Mutex::new(Some(f));

    ctrlc::set_handler(move || {
        if let Some(f) = call_once.lock().unwrap().take() {
            println!("Starting graceful shutdown");
            f();
        } else {
            println!("Already sent signal to start graceful shutdown");
        }
    })
    .unwrap();
}
<commit_msg>Use `HEROKU_SLUG_COMMIT` if available<commit_after>#![warn(clippy::all, rust_2018_idioms)]
#![allow(clippy::unknown_clippy_lints)]

use cargo_registry::{boot, App, Env};
use std::{
    borrow::Cow,
    fs::File,
    sync::{mpsc::channel, Arc, Mutex},
    thread,
    time::Duration,
};

use civet::Server as CivetServer;
use conduit_hyper::Service;
use futures_util::future::FutureExt;
use reqwest::blocking::Client;
use sentry::{ClientOptions, IntoDsn};

const CORE_THREADS: usize = 4;

#[allow(clippy::large_enum_variant)]
enum Server {
    Civet(CivetServer),
    Hyper(tokio::runtime::Runtime, tokio::task::JoinHandle<()>),
}

use Server::*;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let _sentry = dotenv::var("SENTRY_DSN_API")
        .ok()
        .into_dsn()
        .expect("SENTRY_DSN_API is not a valid Sentry DSN value")
        .map(|dsn| {
            let mut opts = ClientOptions::from(dsn);
            opts.environment = Some(
                dotenv::var("SENTRY_ENV_API")
                    .map(Cow::Owned)
                    .expect("SENTRY_ENV_API must be set when using SENTRY_DSN_API"),
            );

            if let Some(commit) = option_env!("CRATES_IO_GIT_COMMIT") {
                opts.release = Some(commit.into());
            } else if let Ok(commit) = dotenv::var("HEROKU_SLUG_COMMIT") {
                opts.release = Some(commit.into());
            }

            sentry::init(opts)
        });

    // Initialize logging
    env_logger::init();

    let config = cargo_registry::Config::default();
    let client = Client::new();

    let app = App::new(config.clone(), Some(client));
    let app = cargo_registry::build_handler(Arc::new(app));

    // On every server restart, ensure the categories available in the database match
    // the information in *src/categories.toml*.
    let categories_toml = include_str!("../boot/categories.toml");
    boot::categories::sync(categories_toml).unwrap();

    let heroku = dotenv::var("HEROKU").is_ok();
    let fastboot = dotenv::var("USE_FASTBOOT").is_ok();

    let port = if heroku {
        8888
    } else {
        dotenv::var("PORT")
            .ok()
            .and_then(|s| s.parse().ok())
            .unwrap_or(8888)
    };
    let threads = dotenv::var("SERVER_THREADS")
        .map(|s| s.parse().expect("SERVER_THREADS was not a valid number"))
        .unwrap_or_else(|_| {
            if config.env == Env::Development {
                5
            } else {
                // A large default because this can be easily changed via env and in production we
                // want the logging middleware to accurately record the start time.
                500
            }
        });

    let server = if dotenv::var("WEB_USE_CIVET").is_err() {
        use tokio::io::AsyncWriteExt;
        use tokio::signal::unix::{signal, SignalKind};

        println!("Booting with a hyper based server");

        let mut rt = tokio::runtime::Builder::new()
            .threaded_scheduler()
            .enable_all()
            .core_threads(CORE_THREADS)
            .max_threads(threads as usize + CORE_THREADS)
            .build()
            .unwrap();

        let handler = Arc::new(conduit_hyper::BlockingHandler::new(app));
        let make_service =
            hyper::service::make_service_fn(move |socket: &hyper::server::conn::AddrStream| {
                let addr = socket.remote_addr();
                let handler = handler.clone();
                async move { Service::from_blocking(handler, addr) }
            });

        let addr = ([127, 0, 0, 1], port).into();
        #[allow(clippy::async_yields_async)]
        let server = rt.block_on(async { hyper::Server::bind(&addr).serve(make_service) });

        let mut sig_int = rt.block_on(async { signal(SignalKind::interrupt()) })?;
        let mut sig_term = rt.block_on(async { signal(SignalKind::terminate()) })?;

        let server = server.with_graceful_shutdown(async move {
            // Wait for either signal
            futures_util::select! {
                _ = sig_int.recv().fuse() => (),
                _ = sig_term.recv().fuse() => (),
            };
            let mut stdout = tokio::io::stdout();
            stdout.write_all(b"Starting graceful shutdown\n").await.ok();
        });

        let server = rt.spawn(async { server.await.unwrap() });
        Hyper(rt, server)
    } else {
        println!("Booting with a civet based server");
        let mut cfg = civet::Config::new();
        cfg.port(port).threads(threads).keep_alive(true);
        Civet(CivetServer::start(cfg, app).unwrap())
    };

    println!("listening on port {}", port);

    // Give tokio a chance to spawn the first worker thread
    thread::sleep(Duration::from_millis(10));

    // Creating this file tells heroku to tell nginx that the application is ready
    // to receive traffic.
    if heroku {
        let path = if fastboot {
            "/tmp/backend-initialized"
        } else {
            "/tmp/app-initialized"
        };
        println!("Writing to {}", path);
        File::create(path).unwrap();
    }

    // Block the main thread until the server has shutdown
    match server {
        Hyper(mut rt, server) => {
            rt.block_on(async { server.await.unwrap() });
        }
        Civet(server) => {
            let (tx, rx) = channel::<()>();
            ctrlc_handler(move || tx.send(()).unwrap_or(()));
            rx.recv().unwrap();
            drop(server);
        }
    }

    println!("Server has gracefully shutdown!");
    Ok(())
}

fn ctrlc_handler<F>(f: F)
where
    F: FnOnce() + Send + 'static,
{
    let call_once = Mutex::new(Some(f));

    ctrlc::set_handler(move || {
        if let Some(f) = call_once.lock().unwrap().take() {
            println!("Starting graceful shutdown");
            f();
        } else {
            println!("Already sent signal to start graceful shutdown");
        }
    })
    .unwrap();
}
<|endoftext|>


<commit_before>use std::collections::{Bitv, VecMap};
use std::num::Int;
use std::{os, mem};

use gc_object::GCObject;
use constants::{BLOCK_SIZE, LINE_SIZE, NUM_LINES_PER_BLOCK};

pub struct BlockInfo {
    mmap: os::MemoryMap,
    line_counter: VecMap<u8>,
    line_map: Bitv,
}

impl BlockInfo {
    pub fn new(mmap: os::MemoryMap) -> BlockInfo {
        debug_assert!(mmap.len() > mem::size_of::<BlockInfo>());
        return BlockInfo {
            mmap: mmap,
            line_counter: VecMap::with_capacity(NUM_LINES_PER_BLOCK),
            line_map: Bitv::from_elem(NUM_LINES_PER_BLOCK, false),
        }
    }

    pub fn into_memory_map(self) -> os::MemoryMap {
        return self.mmap;
    }

    pub fn scan_block(&self, last_high_offset: u16) -> Option<(u16, u16)> {
        debug!("Scanning block {:p} for a hole", self);
        let last_high_index = last_high_offset as uint / LINE_SIZE;
        let mut low_index = NUM_LINES_PER_BLOCK - 1;
        for index in range(last_high_index + 1, NUM_LINES_PER_BLOCK) {
            if self.line_counter.get(&index).map_or(true, |c| *c == 0) {
                // +1 to skip the next line in case an object straddles lines
                low_index = index + 1;
                break;
            }
        }
        debug!("Found low index {} in block {:p}", low_index, self);
        let mut high_index = NUM_LINES_PER_BLOCK;
        for index in range(low_index + 1, NUM_LINES_PER_BLOCK) {
            if self.line_counter.get(&index).map_or(false, |c| *c != 0) {
                high_index = index;
                break;
            }
        }
        debug!("Found high index {} in block {:p}", high_index, self);
        return if low_index < (NUM_LINES_PER_BLOCK - 1) {
            Some(((low_index * LINE_SIZE) as u16,
            (high_index * LINE_SIZE - 1) as u16))
        } else {
            debug!("Found no hole in block {:p}", self);
            None
        };
    }

    pub fn offset(&mut self, offset: uint) -> *mut GCObject {
        return unsafe{ self.mmap.data().offset(offset as int) } as *mut GCObject;
    }

    fn object_to_line_num(object: *const GCObject) -> uint {
        return (object as uint % BLOCK_SIZE) / LINE_SIZE;
    }

    pub fn set_line_mark(&mut self, object: *const GCObject, mark: bool) {
        let line_num = BlockInfo::object_to_line_num(object);
        self.line_map.set(line_num, mark);
        debug!("Set line {} in block {:p}", mark, self);
    }

    pub fn get_line_mark(&self, object: *const GCObject) -> bool {
        return self.line_map
                   .get(BlockInfo::object_to_line_num(object))
                   .unwrap_or(false);
    }

    fn update_line_nums<F>(&mut self, object: *const GCObject, newval: u8,
                           f: F) where F: Fn(u8, u8) -> u8 {
        // This calculates how many lines are affected starting from a
        // LINE_SIZE aligned address. So it might not mark enough lines. But
        // that does not matter as we always skip a line in scan_block()
        let line_num = BlockInfo::object_to_line_num(object);
        let object_size = unsafe{ (*object).object_size() };
        for line in range(line_num, line_num + (object_size / LINE_SIZE) + 1) {
            self.line_counter.update(line, newval, |x, y| f(x, y));
            debug!("Change line count for line {} to {}", line,
                   self.line_counter.get(&line).unwrap());
        }
    }

    pub fn increment_lines(&mut self, object: *const GCObject) {
        debug!("Increment lines in block {:p}", self);
        self.update_line_nums(object, 1, |old, new| old + new);
    }

    pub fn decrement_lines(&mut self, object: *const GCObject) {
        debug!("Decrement lines in block {:p}", self);
        self.update_line_nums(object, 0, |old, _| Int::saturating_sub(old, 1));
    }

    pub fn clear_line_counts(&mut self) {
        self.line_counter.clear();
        self.line_map.clear();
    }

    pub fn is_empty(&self) -> bool {
        return self.line_counter.values().all(|v| *v == 0);
    }
}

<commit_msg>Optimize BlockInfo.update_line_nums()<commit_after>use std::collections::{Bitv, VecMap};
use std::num::Int;
use std::{os, mem};

use gc_object::GCObject;
use constants::{BLOCK_SIZE, LINE_SIZE, NUM_LINES_PER_BLOCK};

pub struct BlockInfo {
    mmap: os::MemoryMap,
    line_counter: VecMap<u8>,
    line_map: Bitv,
}

impl BlockInfo {
    pub fn new(mmap: os::MemoryMap) -> BlockInfo {
        debug_assert!(mmap.len() > mem::size_of::<BlockInfo>());
        return BlockInfo {
            mmap: mmap,
            line_counter: VecMap::with_capacity(NUM_LINES_PER_BLOCK),
            line_map: Bitv::from_elem(NUM_LINES_PER_BLOCK, false),
        }
    }

    pub fn into_memory_map(self) -> os::MemoryMap {
        return self.mmap;
    }

    pub fn scan_block(&self, last_high_offset: u16) -> Option<(u16, u16)> {
        debug!("Scanning block {:p} for a hole", self);
        let last_high_index = last_high_offset as uint / LINE_SIZE;
        let mut low_index = NUM_LINES_PER_BLOCK - 1;
        for index in range(last_high_index + 1, NUM_LINES_PER_BLOCK) {
            if self.line_counter.get(&index).map_or(true, |c| *c == 0) {
                // +1 to skip the next line in case an object straddles lines
                low_index = index + 1;
                break;
            }
        }
        debug!("Found low index {} in block {:p}", low_index, self);
        let mut high_index = NUM_LINES_PER_BLOCK;
        for index in range(low_index + 1, NUM_LINES_PER_BLOCK) {
            if self.line_counter.get(&index).map_or(false, |c| *c != 0) {
                high_index = index;
                break;
            }
        }
        debug!("Found high index {} in block {:p}", high_index, self);
        return if low_index < (NUM_LINES_PER_BLOCK - 1) {
            Some(((low_index * LINE_SIZE) as u16,
            (high_index * LINE_SIZE - 1) as u16))
        } else {
            debug!("Found no hole in block {:p}", self);
            None
        };
    }

    pub fn offset(&mut self, offset: uint) -> *mut GCObject {
        return unsafe{ self.mmap.data().offset(offset as int) } as *mut GCObject;
    }

    fn object_to_line_num(object: *const GCObject) -> uint {
        return (object as uint % BLOCK_SIZE) / LINE_SIZE;
    }

    pub fn set_line_mark(&mut self, object: *const GCObject, mark: bool) {
        let line_num = BlockInfo::object_to_line_num(object);
        self.line_map.set(line_num, mark);
        debug!("Set line {} in block {:p}", mark, self);
    }

    pub fn get_line_mark(&self, object: *const GCObject) -> bool {
        return self.line_map
                   .get(BlockInfo::object_to_line_num(object))
                   .unwrap_or(false);
    }

    fn update_line_nums(&mut self, object: *const GCObject, increment: bool) {
        // This calculates how many lines are affected starting from a
        // LINE_SIZE aligned address. So it might not mark enough lines. But
        // that does not matter as we always skip a line in scan_block()
        let line_num = BlockInfo::object_to_line_num(object);
        let object_size = unsafe{ (*object).object_size() };
        for line in range(line_num, line_num + (object_size / LINE_SIZE) + 1) {
            match increment {
                true => {
                    self.line_counter.update(line, 1, |o, n| o + n);
                    debug!("Incremented line count for line {} to {}", line,
                           self.line_counter.get(&line).unwrap());
                },
                false => {
                    self.line_counter.update(line, 0,
                                             |o, _| Int::saturating_sub(o, 1));
                    debug!("Decremented line count for line {} to {}", line,
                           self.line_counter.get(&line).unwrap());
                }
            }
        }
    }

    pub fn increment_lines(&mut self, object: *const GCObject) {
        self.update_line_nums(object, true);
    }

    pub fn decrement_lines(&mut self, object: *const GCObject) {
        self.update_line_nums(object, false);
    }

    pub fn clear_line_counts(&mut self) {
        self.line_counter.clear();
        self.line_map.clear();
    }

    pub fn is_empty(&self) -> bool {
        return self.line_counter.values().all(|v| *v == 0);
    }
}

<|endoftext|>


<commit_before>//! A splittable pseudorandom generator based on the [Chaskey
//! MAC](http://mouha.be/chaskey/).  **This is not intended to be a
//! cryptographically secure PRNG.**
//!
//! Like `SipRng`, this is broadly modeled after Claessen and Pałka's
//! splittable PRNGs, but with a different choice of cryptographic
//! primitive.
//!
//! ## References
//!
//! * Mouha, Nicky, Bart Mennik, Anthony Van Herrewege, Dai Watanabe,
//!   Bart Preneet and Ingrid Verbauwhede.  2014.  ["Chaskey: An
//!   Efficient MAC Algorithm for 32-bit
//!   Microcontrollers."](https://eprint.iacr.org/2014/386.pdf)
//!   Cryptology ePrint Archive, Report 2014/386.

//! * Claessen, Koen and Michał H. Pałka.  2013.  ["Splittable
//!   Pseudorandom Number Generators using Cryptographic
//!   Hashing."](http://publications.lib.chalmers.se/records/fulltext/183348/local_183348.pdf)
//!   *Haskell '13: Proceedings of the 2013 ACM SIGPLAN symposium on
//!   Haskell*, pp. 47-58.

use rand::{Rand, Rng, SeedableRng};
use super::{SplitRng, SplitPrf};
use std::u32;

/// A splittable pseudorandom generator based on Chaskey.
pub struct ChaskeyRng {
    // The state of the splittable RNG, properly speaking.
      v: [u32; 4],
     k1: [u32; 4],
    ctr: u64,

    // We buffer the 128-bit raw outputs of the RNG to speed it up a bit.
    buf: [u32; 4],
    i: usize
}

/// A PRF taken off a `ChaskeyRng`.
pub struct ChaskeyPrf(ChaskeyRng);

impl ChaskeyRng {
    pub fn new(seed: [u32; 4]) -> ChaskeyRng {
        let mut result = ChaskeyRng {
              v: seed,
             k1: times_two(seed),
            ctr: 0,
            buf: [0u32; 4],
              i: 0
        };
        result.advance();
        result
    }

    fn reseed(&mut self, seed: [u32; 4]) {
        self.v = seed;
        self.k1 = times_two(seed);
        self.ctr = 0;
        self.buf = [0u32; 4];
        self.i = 0;
        self.advance();
    }

    fn clone(&self) -> ChaskeyRng {
        ChaskeyRng {
              v: self.v,
             k1: self.k1,
            ctr: self.ctr,
            buf: self.buf,
              i: self.i
        }
    }

    #[inline]
    fn advance(&mut self) {
        let block = [0, 0, lsb32(self.ctr), msb32(self.ctr)];
        xor_u32x4(&mut self.buf, &block);
        xor_u32x4(&mut self.buf, &self.k1);
        permute8(&mut self.buf);
        xor_u32x4(&mut self.buf, &self.k1);

        self.ctr = self.ctr.wrapping_add(1);
        self.i = 0;
    }

    #[inline]
    fn descend(&mut self, i: u32) {
        let block = [u32::MAX, i, lsb32(self.ctr), msb32(self.ctr)];
        xor_u32x4(&mut self.buf, &block);
        permute8(&mut self.v);

        self.ctr = 0;
        self.advance();
    }

}

impl SplitPrf<ChaskeyRng> for ChaskeyPrf {
    fn call(&self, i: u32) -> ChaskeyRng {
        let mut r = self.0.clone();
        r.descend(i);
        r
    }
}

impl SplitRng for ChaskeyRng {
    type Prf = ChaskeyPrf;

    fn split(&mut self) -> Self {
        let mut child = self.clone();
        self.descend(0);
        child.descend(1);
        child
    }

    fn splitn(&mut self) -> ChaskeyPrf {
        ChaskeyPrf(self.split())
    }

}

impl Rng for ChaskeyRng {
    #[inline]
    fn next_u32(&mut self) -> u32 {
        if self.i >= 4 {
            self.advance();
        }
        let result = self.buf[self.i];
        self.i += 1;
        result
    }
}

impl SeedableRng<[u32; 4]> for ChaskeyRng {

    fn reseed(&mut self, seed: [u32; 4]) {
        ChaskeyRng::reseed(self, seed);
    }

    fn from_seed(seed: [u32; 4]) -> ChaskeyRng {
        ChaskeyRng::new(seed)
    }
}

impl Rand for ChaskeyRng {
    fn rand<R: Rng>(other: &mut R) -> ChaskeyRng {
        ChaskeyRng::new(other.gen::<[u32; 4]>())
    }
}

/*
 * Chaskey's building blocks.
 */

/// Function used in the Chaskey key schedule.
#[inline(always)]
pub fn times_two(key: [u32; 4]) -> [u32; 4] {
    const C: [u32; 2] = [0x00, 0x87];
    [key[0].wrapping_shl(1) ^ C[key[3].wrapping_shr(31) as usize],
     key[1].wrapping_shl(1) ^ key[0].wrapping_shr(31),
     key[2].wrapping_shl(1) ^ key[1].wrapping_shr(31),
     key[3].wrapping_shl(1) ^ key[2].wrapping_shr(31)]
}

/// XOR a `[u32; 4]` value into the Chaskey state.
#[inline(always)]
pub fn xor_u32x4(state: &mut [u32; 4], block: &[u32; 4]) {
    state[0] ^= block[0];
    state[1] ^= block[1];
    state[2] ^= block[2];
    state[3] ^= block[3];
}

// The original Chaskey permutation (8 rounds).
#[inline(always)]
pub fn permute8(state: &mut [u32; 4]) {
    permute4(state); permute4(state);
}

#[inline(always)]
fn permute4(state: &mut [u32; 4]) {
    round(state); round(state);
    round(state); round(state);
}

/// The Chaskey round function.
#[inline(always)]
pub fn round(v: &mut [u32; 4]) {
    v[0]  = v[0].wrapping_add(v[1]); v[2]  = v[2].wrapping_add(v[3]);
    v[1]  = v[1].rotate_left(5);     v[3]  = v[3].rotate_left(8);
    v[1] ^= v[0];                    v[3] ^= v[2];
    v[0]  = v[0].rotate_left(16);

    v[2]  = v[2].wrapping_add(v[1]); v[0]  = v[0].wrapping_add(v[3]);
    v[1]  = v[1].rotate_left(7);     v[3]  = v[3].rotate_left(13);
    v[1] ^= v[2];                    v[3] ^= v[0];
    v[2]  = v[2].rotate_left(16);
}

/*
 * Utility functions.
 */

/// Take the least 32 significant bits of an `u64`.
#[inline(always)]
fn lsb32(n: u64) -> u32 {
    const MASK: u64 = 0xffff_ffff;
    (n & MASK) as u32
}

/// Take the most 32 significant bits of an `u64`.
#[inline(always)]
fn msb32(n: u64) -> u32 {
    lsb32(n.wrapping_shr(32))
}

#[cfg(test)]
mod tests {
    use rand::Rng;
    use rand::os::OsRng;
    use chaskeyrng::ChaskeyRng;

    fn gen_chaskeyrng() -> ChaskeyRng {
        let mut osrng = OsRng::new().ok().expect("Could not create OsRng");
        osrng.gen()
    }

    #[test]
    fn test_split_rand_independence() {
        ::tests::test_split_rand_independence(&mut gen_chaskeyrng());
    }

    #[test]
    fn test_split_rand_closure() {
        ::tests::test_split_rand_closure(&mut gen_chaskeyrng());
    }

    #[test]
    fn test_split_rand_split() {
        ::tests::test_split_rand_split(&mut gen_chaskeyrng());
    }

    fn gen_seed() -> [u32; 4] {
        let mut osrng = OsRng::new().ok().expect("Could not create OsRng");
        osrng.gen()
    }

    #[test]
    fn test_rng_rand_seeded() {
        let seed = gen_seed();
        ::tests::test_rng_rand_seeded::<ChaskeyRng, [u32; 4]>(seed);
    }

    #[test]
    fn test_rng_seeded() {
        let seed = gen_seed();
        ::tests::test_rng_seeded::<ChaskeyRng, [u32; 4]>(seed);
    }

    #[test]
    fn test_rng_reseed() {
        let seed = gen_seed();
        ::tests::test_rng_reseed::<ChaskeyRng, [u32; 4]>(seed);
    }
}
<commit_msg>`ChaskeyRng` is so easy to get wrong.  Fix some issues and comment profusely the tricky parts.<commit_after>//! A splittable pseudorandom generator based on the [Chaskey
//! MAC](http://mouha.be/chaskey/).  **This is not intended to be a
//! cryptographically secure PRNG.**
//!
//! Like `SipRng`, this is broadly modeled after Claessen and Pałka's
//! splittable PRNGs, but with a different choice of cryptographic
//! primitive.
//!
//! ## References
//!
//! * Mouha, Nicky, Bart Mennik, Anthony Van Herrewege, Dai Watanabe,
//!   Bart Preneet and Ingrid Verbauwhede.  2014.  ["Chaskey: An
//!   Efficient MAC Algorithm for 32-bit
//!   Microcontrollers."](https://eprint.iacr.org/2014/386.pdf)
//!   Cryptology ePrint Archive, Report 2014/386.

//! * Claessen, Koen and Michał H. Pałka.  2013.  ["Splittable
//!   Pseudorandom Number Generators using Cryptographic
//!   Hashing."](http://publications.lib.chalmers.se/records/fulltext/183348/local_183348.pdf)
//!   *Haskell '13: Proceedings of the 2013 ACM SIGPLAN symposium on
//!   Haskell*, pp. 47-58.

use rand::{Rand, Rng, SeedableRng};
use super::{SplitRng, SplitPrf};
use std::u32;

/// A splittable pseudorandom generator based on Chaskey.
pub struct ChaskeyRng {
    // The state of the splittable RNG, properly speaking.
    state: [u32; 4],
       k1: [u32; 4],
      ctr: u64,

    // We buffer the 128-bit raw outputs of the RNG to speed it up a bit.
      buf: [u32; 4],
        i: usize
}

/// A PRF taken off a `ChaskeyRng`.
pub struct ChaskeyPrf(ChaskeyRng);

impl ChaskeyRng {
    pub fn new(seed: [u32; 4]) -> ChaskeyRng {
        let mut result = ChaskeyRng {
            state: seed,
               k1: times_two(seed),
              ctr: 0,

            // We keep a buffer of the most recent words we generated
            // so that we don't call the MAC as often.
              buf: [0u32; 4],
            // Current position within `buf`.
                i: 0
        };
        result.advance();
        result
    }

    fn reseed(&mut self, seed: [u32; 4]) {
        self.state = seed;
        self.k1    = times_two(seed);
        self.ctr   = 0;
        self.buf   = [0u32; 4];
        self.i     = 0;
        self.advance();
    }

    fn clone(&self) -> ChaskeyRng {
        ChaskeyRng {
            state: self.state,
               k1: self.k1,
              ctr: self.ctr,
              buf: self.buf,
                i: self.i
        }
    }

    #[inline]
    fn advance(&mut self) {
        // TRICKY CODE: We do this in the `buf` of the generator, so
        // as to leave the `state` untouched.  Effectively what we're
        // doing here is appending the counter value to a prefix we've
        // processed already, and finishing the Chaskey computation.
        // So we must leave `state` untouched!

        // Copy the `state` into the `buf`, and mix in the counter
        // value.
        self.buf[0] = self.state[0];
        self.buf[1] = self.state[1];
        self.buf[2] = self.state[2] ^ lsb32(self.ctr);
        self.buf[3] = self.state[3] ^ msb32(self.ctr);

        // Finalize the Chaskey computation on `buf`.
        xor_u32x4(&mut self.buf, &self.k1);
        permute8(&mut self.buf);
        xor_u32x4(&mut self.buf, &self.k1);

        self.ctr = self.ctr.wrapping_add(1);
        self.i = 0;
    }

    #[inline]
    fn descend(&mut self, i: u32) {
        self.state[0] ^= u32::MAX;
        self.state[1] ^= i;
        self.state[2] ^= lsb32(self.ctr);
        self.state[3] ^= msb32(self.ctr);
        permute8(&mut self.state);

        self.ctr = 0;
        self.advance();
    }

}

impl SplitPrf<ChaskeyRng> for ChaskeyPrf {
    fn call(&self, i: u32) -> ChaskeyRng {
        let mut r = self.0.clone();
        r.descend(i);
        r
    }
}

impl SplitRng for ChaskeyRng {
    type Prf = ChaskeyPrf;

    fn split(&mut self) -> Self {
        let mut child = self.clone();
        self.descend(0);
        child.descend(1);
        child
    }

    fn splitn(&mut self) -> ChaskeyPrf {
        ChaskeyPrf(self.split())
    }

}

impl Rng for ChaskeyRng {
    #[inline]
    fn next_u32(&mut self) -> u32 {
        if self.i >= 4 {
            self.advance();
        }
        let result = self.buf[self.i];
        self.i += 1;
        result
    }
}

impl SeedableRng<[u32; 4]> for ChaskeyRng {

    fn reseed(&mut self, seed: [u32; 4]) {
        ChaskeyRng::reseed(self, seed);
    }

    fn from_seed(seed: [u32; 4]) -> ChaskeyRng {
        ChaskeyRng::new(seed)
    }
}

impl Rand for ChaskeyRng {
    fn rand<R: Rng>(other: &mut R) -> ChaskeyRng {
        ChaskeyRng::new(other.gen::<[u32; 4]>())
    }
}

/*
 * Chaskey's building blocks.
 */

/// Function used in the Chaskey key schedule.
#[inline(always)]
pub fn times_two(key: [u32; 4]) -> [u32; 4] {
    const C: [u32; 2] = [0x00, 0x87];
    [key[0].wrapping_shl(1) ^ C[key[3].wrapping_shr(31) as usize],
     key[1].wrapping_shl(1) ^ key[0].wrapping_shr(31),
     key[2].wrapping_shl(1) ^ key[1].wrapping_shr(31),
     key[3].wrapping_shl(1) ^ key[2].wrapping_shr(31)]
}

/// XOR a `[u32; 4]` value into the Chaskey state.
#[inline(always)]
pub fn xor_u32x4(state: &mut [u32; 4], block: &[u32; 4]) {
    state[0] ^= block[0];
    state[1] ^= block[1];
    state[2] ^= block[2];
    state[3] ^= block[3];
}

// The original Chaskey permutation (8 rounds).
#[inline(always)]
pub fn permute8(state: &mut [u32; 4]) {
    permute4(state); permute4(state);
}

#[inline(always)]
fn permute4(state: &mut [u32; 4]) {
    round(state); round(state);
    round(state); round(state);
}

/// The Chaskey round function.
#[inline(always)]
pub fn round(v: &mut [u32; 4]) {
    v[0]  = v[0].wrapping_add(v[1]); v[2]  = v[2].wrapping_add(v[3]);
    v[1]  = v[1].rotate_left(5);     v[3]  = v[3].rotate_left(8);
    v[1] ^= v[0];                    v[3] ^= v[2];
    v[0]  = v[0].rotate_left(16);

    v[2]  = v[2].wrapping_add(v[1]); v[0]  = v[0].wrapping_add(v[3]);
    v[1]  = v[1].rotate_left(7);     v[3]  = v[3].rotate_left(13);
    v[1] ^= v[2];                    v[3] ^= v[0];
    v[2]  = v[2].rotate_left(16);
}

/*
 * Utility functions.
 */

/// Take the least 32 significant bits of an `u64`.
#[inline(always)]
fn lsb32(n: u64) -> u32 {
    const MASK: u64 = 0xffff_ffff;
    (n & MASK) as u32
}

/// Take the most 32 significant bits of an `u64`.
#[inline(always)]
fn msb32(n: u64) -> u32 {
    lsb32(n.wrapping_shr(32))
}

#[cfg(test)]
mod tests {
    use rand::Rng;
    use rand::os::OsRng;
    use chaskeyrng::ChaskeyRng;

    fn gen_chaskeyrng() -> ChaskeyRng {
        let mut osrng = OsRng::new().ok().expect("Could not create OsRng");
        osrng.gen()
    }

    #[test]
    fn test_split_rand_independence() {
        ::tests::test_split_rand_independence(&mut gen_chaskeyrng());
    }

    #[test]
    fn test_split_rand_closure() {
        ::tests::test_split_rand_closure(&mut gen_chaskeyrng());
    }

    #[test]
    fn test_split_rand_split() {
        ::tests::test_split_rand_split(&mut gen_chaskeyrng());
    }

    fn gen_seed() -> [u32; 4] {
        let mut osrng = OsRng::new().ok().expect("Could not create OsRng");
        osrng.gen()
    }

    #[test]
    fn test_rng_rand_seeded() {
        let seed = gen_seed();
        ::tests::test_rng_rand_seeded::<ChaskeyRng, [u32; 4]>(seed);
    }

    #[test]
    fn test_rng_seeded() {
        let seed = gen_seed();
        ::tests::test_rng_seeded::<ChaskeyRng, [u32; 4]>(seed);
    }

    #[test]
    fn test_rng_reseed() {
        let seed = gen_seed();
        ::tests::test_rng_reseed::<ChaskeyRng, [u32; 4]>(seed);
    }
}
<|endoftext|>