<commit_before>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! A [Fleetspeak] client connector library.
//!
//! This library exposes a set of functions for writing client-side Fleetspeak
//! services. Each of these functions operates on a global connection object
//! that is lazily established. If this global connection cannot be established,
//! the library will panic (because without this connection Fleetspeak will shut
//! the service down anyway).
//!
//! Note that each service should send startup information upon its inception
//! and continue to heartbeat from time to time to notify the Fleetspeak client
//! that it did not get stuck.
//!
//! [Fleetspeak]: https://github.com/google/fleetspeak

mod connection;
mod error;

use std::fs::File;
use std::sync::Mutex;
use std::time::Duration;

use lazy_static::lazy_static;

pub use self::connection::Packet;
pub use self::error::{ReadError, WriteError};

/// Sends a heartbeat signal to the Fleetspeak client.
///
/// All client services should heartbeat from time to time. Otherwise, from the
/// Fleetspeak perspective, the service is unresponsive and should be restarted.
///
/// The exact frequency of the required heartbeat is defined in the service
/// configuration file.
pub fn heartbeat() -> Result<(), WriteError> {
    locked(&CONNECTION.output, |buf| self::connection::heartbeat(buf))
}

/// Sends a system message with startup information to the Fleetspeak client.
///
/// All clients are required to send this information on startup. If the client
/// does not receive this information quickly enough, the service will be
/// killed.
///
/// The `version` string should contain a self-reported version of the service.
/// This data is used primarily for statistics.
pub fn startup(version: &str) -> Result<(), WriteError> {
    locked(&CONNECTION.output, |buf| self::connection::startup(buf, version))
}

/// Sends the message to the Fleetspeak server.
///
/// The message is delivered to the server-side service as specified by the
/// packet and optionally tagged with a type if specified. This optional message
/// type is irrelevant for Fleetspeak but might be useful for the service the
/// message is delivered to.
///
/// In case of any I/O failure or malformed message (e.g. due to encoding
/// problems), an error is reported.
///
/// # Examples
///
/// ```no_run
/// use fleetspeak::Packet;
///
/// fleetspeak::send(Packet {
///     service: String::from("example"),
///     kind: None,
///     data: String::from("Hello, World!"),
/// }).expect("failed to send the packet");
/// ```
pub fn send<M>(packet: Packet<M>) -> Result<(), WriteError>
where
    M: prost::Message,
{
    locked(&CONNECTION.output, |buf| self::connection::send(buf, packet))
}

/// Receives a message from the Fleetspeak server.
///
/// This function will block until there is a message to be read from the input.
/// Note that in particular it means your service will be unable to heartbeat
/// properly. If you are not expecting the message to arrive quickly, you should
/// use [`collect`] instead.
///
/// In case of any I/O failure or malformed message (e.g. due to parsing issues
/// or when some fields are not being present), an error is reported.
///
/// [`collect`]: fn.collect.html
///
/// # Examples
///
/// ```no_run
/// match fleetspeak::receive::<String>() {
///     Ok(packet) => println!("Hello, {}!", packet.data),
///     Err(error) => eprintln!("failed to receive the packet: {}", error),
/// }
/// ```
pub fn receive<M>() -> Result<Packet<M>, ReadError>
where
    M: prost::Message + Default,
{
    locked(&CONNECTION.input, |buf| self::connection::receive(buf))
}

/// Collects a message from the Fleetspeak server.
///
/// Unlike [`receive`], `collect` will send heartbeat signals while polling for
/// the message at the specified `rate`.
///
/// This function is useful in the main loop of your service when it is not
/// supposed to do anything until a request from the server arrives. If your
/// service is actually awaiting for a specific message to come, you should
/// use [`receive`] instead.
///
/// In case of any I/O failure or malformed message (e.g. due to parsing issues
/// or when some fields are not being present), an error is reported.
///
/// [`receive`]: fn.receive.html
///
/// # Examples
///
/// ```no_run
/// use std::time::Duration;
///
/// match fleetspeak::collect::<String>(Duration::from_secs(1)) {
///     Ok(packet) => println!("Hello, {}!", packet.data),
///     Err(error) => eprintln!("failed to collected the packet: {}", error),
/// }
/// ```
pub fn collect<M>(rate: Duration) -> Result<Packet<M>, std::io::Error>
where
    M: prost::Message + Default + 'static,
{
    // TODO: Refactor this code once `!` stabilizes.
    let (sender, receiver) = std::sync::mpsc::channel();

    std::thread::spawn(move || {
        loop {
            use std::sync::mpsc::TryRecvError::*;

            // The heartbeat thread should stop itself when it receives a signal
            // to do so (or when the channel is closed). Otherwise, it should
            // keep heartbeating.
            match receiver.try_recv() {
                Ok(()) => return,
                Err(Empty) => (),
                Err(Disconnected) => return,
            }

            // Ignoring heartbeat errors is not great, but they can occur only
            // in very rare cases and any subsequent write operations are going
            // to fail soon anyway. Hence, we drop the error on the floor and
            // shut the thread down, hoping that the main thread will notice the
            // problem as soon as it tries to write something. In case the main
            // thread blocks indefinitely, Fleetspeak should figure out that the
            // service is unresponsive and kill it eventually.
            // TODO: Add logging.
            match heartbeat() {
                Ok(()) => (),
                Err(_) => return,
            }

            std::thread::sleep(rate);
        }
    });

    let packet = receive()?;

    // Notify the heartbeat thread to shut down. We do not really care whether
    // the thread was aborted (it should never be, but if it was, we just carry
    // on).
    // TODO: Add logging.
    let _ = sender.send(());

    Ok(packet)
}

/// A connection to the Fleetspeak client.
///
/// The connection is realized through two files (specified by descriptors given
/// by the Fleetspeak client as environment variables): input and output. Each
/// of these files is guarded by a separate mutex to allow writing (e.g. for
/// sending heartbeat signals) when another thread might be busy with reading
/// messages.
struct Connection {
    input: Mutex<File>,
    output: Mutex<File>,
}

lazy_static! {
    static ref CONNECTION: Connection = {
        let mut input = open("FLEETSPEAK_COMMS_CHANNEL_INFD");
        let mut output = open("FLEETSPEAK_COMMS_CHANNEL_OUTFD");

        use self::connection::handshake;
        handshake(&mut input, &mut output).expect("handshake failure");

        Connection {
            input: Mutex::new(input),
            output: Mutex::new(output),
        }
    };
}

/// Executes the given function with a file extracted from the mutex.
///
/// It might happen that the mutex becomes poisoned and this call will panic in
/// result. This should not be a problem in practice, because mutex poisoning
/// is a result of one of the threads being aborted. In case of a such scenario,
/// it is likely the service needs to be restarted anyway.
fn locked<F, T, E>(mutex: &Mutex<File>, f: F) -> Result<T, E>
where
    F: FnOnce(&mut File) -> Result<T, E>
{
    let mut file = mutex.lock().expect("poisoned connection mutex");
    f(&mut file)
}

/// Opens a file object pointed by an environment variable.
///
/// Note that this function will panic if the environment variable `var` is not
/// a valid file descriptor (in which case the library cannot be initialized and
/// the service is unlikely to work anyway).
fn open(var: &str) -> File {
    let fd = std::env::var(var)
        .expect(&format!("invalid variable `{}`", var))
        .parse()
        .expect(&format!("failed to parse file descriptor"));

    // TODO: Add support for Windows.
    unsafe {
        std::os::unix::io::FromRawFd::from_raw_fd(fd)
    }
}
<commit_msg>Clarify the doc comment about the heartbeat rate.<commit_after>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! A [Fleetspeak] client connector library.
//!
//! This library exposes a set of functions for writing client-side Fleetspeak
//! services. Each of these functions operates on a global connection object
//! that is lazily established. If this global connection cannot be established,
//! the library will panic (because without this connection Fleetspeak will shut
//! the service down anyway).
//!
//! Note that each service should send startup information upon its inception
//! and continue to heartbeat from time to time to notify the Fleetspeak client
//! that it did not get stuck.
//!
//! [Fleetspeak]: https://github.com/google/fleetspeak

mod connection;
mod error;

use std::fs::File;
use std::sync::Mutex;
use std::time::Duration;

use lazy_static::lazy_static;

pub use self::connection::Packet;
pub use self::error::{ReadError, WriteError};

/// Sends a heartbeat signal to the Fleetspeak client.
///
/// All client services should heartbeat from time to time. Otherwise, from the
/// Fleetspeak perspective, the service is unresponsive and should be restarted.
///
/// The exact frequency of the required heartbeat is defined in the service
/// configuration file.
pub fn heartbeat() -> Result<(), WriteError> {
    locked(&CONNECTION.output, |buf| self::connection::heartbeat(buf))
}

/// Sends a system message with startup information to the Fleetspeak client.
///
/// All clients are required to send this information on startup. If the client
/// does not receive this information quickly enough, the service will be
/// killed.
///
/// The `version` string should contain a self-reported version of the service.
/// This data is used primarily for statistics.
pub fn startup(version: &str) -> Result<(), WriteError> {
    locked(&CONNECTION.output, |buf| self::connection::startup(buf, version))
}

/// Sends the message to the Fleetspeak server.
///
/// The message is delivered to the server-side service as specified by the
/// packet and optionally tagged with a type if specified. This optional message
/// type is irrelevant for Fleetspeak but might be useful for the service the
/// message is delivered to.
///
/// In case of any I/O failure or malformed message (e.g. due to encoding
/// problems), an error is reported.
///
/// # Examples
///
/// ```no_run
/// use fleetspeak::Packet;
///
/// fleetspeak::send(Packet {
///     service: String::from("example"),
///     kind: None,
///     data: String::from("Hello, World!"),
/// }).expect("failed to send the packet");
/// ```
pub fn send<M>(packet: Packet<M>) -> Result<(), WriteError>
where
    M: prost::Message,
{
    locked(&CONNECTION.output, |buf| self::connection::send(buf, packet))
}

/// Receives a message from the Fleetspeak server.
///
/// This function will block until there is a message to be read from the input.
/// Note that in particular it means your service will be unable to heartbeat
/// properly. If you are not expecting the message to arrive quickly, you should
/// use [`collect`] instead.
///
/// In case of any I/O failure or malformed message (e.g. due to parsing issues
/// or when some fields are not being present), an error is reported.
///
/// [`collect`]: fn.collect.html
///
/// # Examples
///
/// ```no_run
/// match fleetspeak::receive::<String>() {
///     Ok(packet) => println!("Hello, {}!", packet.data),
///     Err(error) => eprintln!("failed to receive the packet: {}", error),
/// }
/// ```
pub fn receive<M>() -> Result<Packet<M>, ReadError>
where
    M: prost::Message + Default,
{
    locked(&CONNECTION.input, |buf| self::connection::receive(buf))
}

/// Collects a message from the Fleetspeak server.
///
/// Unlike [`receive`], `collect` will send heartbeat signals at the specified
/// `rate` while waiting for the message.
///
/// This function is useful in the main loop of your service when it is not
/// supposed to do anything until a request from the server arrives. If your
/// service is actually awaiting for a specific message to come, you should
/// use [`receive`] instead.
///
/// In case of any I/O failure or malformed message (e.g. due to parsing issues
/// or when some fields are not being present), an error is reported.
///
/// [`receive`]: fn.receive.html
///
/// # Examples
///
/// ```no_run
/// use std::time::Duration;
///
/// match fleetspeak::collect::<String>(Duration::from_secs(1)) {
///     Ok(packet) => println!("Hello, {}!", packet.data),
///     Err(error) => eprintln!("failed to collected the packet: {}", error),
/// }
/// ```
pub fn collect<M>(rate: Duration) -> Result<Packet<M>, std::io::Error>
where
    M: prost::Message + Default + 'static,
{
    // TODO: Refactor this code once `!` stabilizes.
    let (sender, receiver) = std::sync::mpsc::channel();

    std::thread::spawn(move || {
        loop {
            use std::sync::mpsc::TryRecvError::*;

            // The heartbeat thread should stop itself when it receives a signal
            // to do so (or when the channel is closed). Otherwise, it should
            // keep heartbeating.
            match receiver.try_recv() {
                Ok(()) => return,
                Err(Empty) => (),
                Err(Disconnected) => return,
            }

            // Ignoring heartbeat errors is not great, but they can occur only
            // in very rare cases and any subsequent write operations are going
            // to fail soon anyway. Hence, we drop the error on the floor and
            // shut the thread down, hoping that the main thread will notice the
            // problem as soon as it tries to write something. In case the main
            // thread blocks indefinitely, Fleetspeak should figure out that the
            // service is unresponsive and kill it eventually.
            // TODO: Add logging.
            match heartbeat() {
                Ok(()) => (),
                Err(_) => return,
            }

            std::thread::sleep(rate);
        }
    });

    let packet = receive()?;

    // Notify the heartbeat thread to shut down. We do not really care whether
    // the thread was aborted (it should never be, but if it was, we just carry
    // on).
    // TODO: Add logging.
    let _ = sender.send(());

    Ok(packet)
}

/// A connection to the Fleetspeak client.
///
/// The connection is realized through two files (specified by descriptors given
/// by the Fleetspeak client as environment variables): input and output. Each
/// of these files is guarded by a separate mutex to allow writing (e.g. for
/// sending heartbeat signals) when another thread might be busy with reading
/// messages.
struct Connection {
    input: Mutex<File>,
    output: Mutex<File>,
}

lazy_static! {
    static ref CONNECTION: Connection = {
        let mut input = open("FLEETSPEAK_COMMS_CHANNEL_INFD");
        let mut output = open("FLEETSPEAK_COMMS_CHANNEL_OUTFD");

        use self::connection::handshake;
        handshake(&mut input, &mut output).expect("handshake failure");

        Connection {
            input: Mutex::new(input),
            output: Mutex::new(output),
        }
    };
}

/// Executes the given function with a file extracted from the mutex.
///
/// It might happen that the mutex becomes poisoned and this call will panic in
/// result. This should not be a problem in practice, because mutex poisoning
/// is a result of one of the threads being aborted. In case of a such scenario,
/// it is likely the service needs to be restarted anyway.
fn locked<F, T, E>(mutex: &Mutex<File>, f: F) -> Result<T, E>
where
    F: FnOnce(&mut File) -> Result<T, E>
{
    let mut file = mutex.lock().expect("poisoned connection mutex");
    f(&mut file)
}

/// Opens a file object pointed by an environment variable.
///
/// Note that this function will panic if the environment variable `var` is not
/// a valid file descriptor (in which case the library cannot be initialized and
/// the service is unlikely to work anyway).
fn open(var: &str) -> File {
    let fd = std::env::var(var)
        .expect(&format!("invalid variable `{}`", var))
        .parse()
        .expect(&format!("failed to parse file descriptor"));

    // TODO: Add support for Windows.
    unsafe {
        std::os::unix::io::FromRawFd::from_raw_fd(fd)
    }
}
<|endoftext|>


<commit_before>use std::ffi::{OsStr, OsString};
use std::io;
use std::ops::{Deref, DerefMut};
use std::os::windows::ffi::{OsStrExt, OsStringExt};
use std::ptr::{null, null_mut};
use std::slice;
use winapi::winrt::hstring::HSTRING;
use winapi::winrt::winstring::*;

pub struct HString {
    __inner: HSTRING,
}

impl HString {
    pub fn new() -> Result<HString, io::Error> {
        Self::with_wide_string(None)
    }

    pub unsafe fn null() -> HString {
        HString {
            __inner: null_mut(),
        }
    }

    pub fn len(&self) -> usize {
        unsafe { WindowsGetStringLen(self.__inner) as usize }
    }

    fn with_wide_string(vec: Option<Vec<u16>>) -> Result<HString, io::Error> {
        let mut handle: HSTRING = null_mut();

        let (ptr, len) = match vec {
            Some(v) => (v.as_ptr(), v.len() as u32),
            None => (null(), 0),
        };

        let ret = unsafe { WindowsCreateString(ptr, len, &mut handle) };

        if ret != 0 {
            return Err(io::Error::last_os_error());
        }

        Ok(HString { __inner: handle })
    }
}

impl Drop for HString {
    fn drop(&mut self) {
        if !self.__inner.is_null() {
            unsafe { WindowsDeleteString(self.__inner) };
        }
    }
}

impl Deref for HString {
    type Target = HSTRING;

    fn deref(&self) -> &HSTRING {
        &self.__inner
    }
}

impl DerefMut for HString {
    fn deref_mut(&mut self) -> &mut HSTRING {
        &mut self.__inner
    }
}

impl From<HString> for String {
    fn from(hstring: HString) -> Self {
        let os_string = OsString::from(hstring);

        os_string
            .as_os_str()
            .to_string_lossy()
            .split('\0')
            .next()
            .unwrap()
            .to_owned()
    }
}

impl From<HString> for OsString {
    fn from(hstring: HString) -> Self {
        let len = hstring.len();
        let buf = unsafe { WindowsGetStringRawBuffer(hstring.__inner, &mut (len as u32)) };
        let s = unsafe { slice::from_raw_parts(buf, len) };
        OsString::from_wide(&s)
    }
}

impl<'a> From<&'a str> for HString {
    fn from(string: &str) -> Self {
        let x = OsStr::new(string).encode_wide().collect();
        HString::with_wide_string(Some(x)).unwrap()
    }
}

impl From<String> for HString {
    fn from(string: String) -> Self {
        HString::from(string.as_str())
    }
}

#[test]
fn test_hstring_from() {
    let t = "This is a test.";
    let h = HString::from(t);
    let s = String::from(h);
    assert_eq!(t, s);
}
<commit_msg>Correct HString construction<commit_after>use std::ffi::{OsStr, OsString};
use std::io;
use std::ops::{Deref, DerefMut};
use std::os::windows::ffi::{OsStrExt, OsStringExt};
use std::ptr::{null, null_mut};
use std::slice;
use winapi::winrt::hstring::HSTRING;
use winapi::winrt::winstring::*;

pub struct HString {
    __inner: HSTRING,
}

impl HString {
    pub fn new() -> Result<HString, io::Error> {
        Self::with_wide_string(Vec::new())
    }

    pub unsafe fn null() -> HString {
        HString {
            __inner: null_mut(),
        }
    }

    pub fn len(&self) -> usize {
        unsafe { WindowsGetStringLen(self.__inner) as usize }
    }

    fn with_wide_string(mut vec: Vec<u16>) -> Result<HString, io::Error> {
        let mut handle: HSTRING = null_mut();

        vec.push(0);
        let (ptr, len) = match vec.len() {
            0 | 1 => (null(), 0),
            n => {
                (vec.as_ptr(), (n - 1) as u32)
            }
        };

        let ret = unsafe { WindowsCreateString(ptr, len, &mut handle) };

        if ret != 0 {
            return Err(io::Error::last_os_error());
        }

        // DO NOT REMOVE THIS, this helps ensure the vec lives longer than the call to WindowsCreateString
        std::mem::drop(vec);

        Ok(HString { __inner: handle })
    }
}

impl Drop for HString {
    fn drop(&mut self) {
        if !self.__inner.is_null() {
            unsafe { WindowsDeleteString(self.__inner) };
        }
    }
}

impl Deref for HString {
    type Target = HSTRING;

    fn deref(&self) -> &HSTRING {
        &self.__inner
    }
}

impl DerefMut for HString {
    fn deref_mut(&mut self) -> &mut HSTRING {
        &mut self.__inner
    }
}

impl From<HString> for String {
    fn from(hstring: HString) -> Self {
        let os_string = OsString::from(hstring);

        os_string
            .as_os_str()
            .to_string_lossy()
            .split('\0')
            .next()
            .unwrap()
            .to_owned()
    }
}

impl From<HString> for OsString {
    fn from(hstring: HString) -> Self {
        let len = hstring.len();
        let buf = unsafe { WindowsGetStringRawBuffer(hstring.__inner, &mut (len as u32)) };
        let s = unsafe { slice::from_raw_parts(buf, len) };
        OsString::from_wide(&s)
    }
}

impl<'a> From<&'a str> for HString {
    fn from(string: &str) -> Self {
        let mut x = OsStr::new(string).encode_wide().collect();
        unsafe { HString::with_wide_string(x).unwrap() }
    }
}

impl From<String> for HString {
    fn from(string: String) -> Self {
        HString::from(string.as_str())
    }
}

#[test]
fn test_hstring_from() {
    let t = "This is a test.";
    let h = HString::from(t);
    let s = String::from(h);
    assert_eq!(t, s);
}
<|endoftext|>


<commit_before>//! Struct and methods for union-find operation.

#![warn(bad_style, missing_docs,
        unused, unused_extern_crates, unused_import_braces,
        unused_qualifications, unused_results, unused_typecasts)]

#![cfg_attr(test, feature(test))]

use std::{mem, usize};
use std::default::Default;
use std::iter::{IntoIterator, FromIterator};

#[derive(Clone, Debug)]
enum UFNode<T> {
    Key(usize),
    Value(T)
}

impl<T> UFNode<T> {
    fn get(&self) -> &T {
        match self {
            &UFNode::Value(ref val) => val,
            &UFNode::Key(_) => panic!()
        }
    }
    fn get_mut(&mut self) -> &mut T {
        match self {
            &mut UFNode::Value(ref mut val) => val,
            &mut UFNode::Key(_) => panic!()
        }
    }
    fn unwrap(self) -> T {
        match self {
            UFNode::Value(val) => val,
            UFNode::Key(_) => panic!()
        }
    }
}

/// The value that can be contained with `UFValue`.
pub trait UFValue: Sized {
    /// Merge two value into one.
    ///
    /// This is used by `UnionFind::union` operation.
    fn merge(lval: Self, rval: Self) -> Merge<Self>;
}

/// Return value of the `UFValue::merege`
#[allow(missing_docs)]
pub enum Merge<T> {
    Left(T), Right(T)
}

/// Reperesents the size of the group.
#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]
pub struct Size(usize);

impl UFValue for Size {
    fn merge(Size(lval): Size, Size(rval): Size) -> Merge<Size> {
        if lval >= rval {
            Merge::Left(Size(lval + rval))
        } else {
            Merge::Right(Size(lval + rval))
        }
    }
}

impl Default for Size {
    fn default() -> Size { Size(1) }
}

/// Struct for union-find operation.
#[derive(Clone, Debug)]
pub struct UnionFind<V = Size> {
    data: Vec<UFNode<V>>
}

impl<T: UFValue = Size> UnionFind<T> {
    /// Creates empty `UnionFind` struct.
    #[inline]
    pub fn new(len: usize) -> UnionFind<T>
        where T: Default
    {
        UnionFind {
            data: (0 .. len)
                .map(|_| UFNode::Value(Default::default()))
                .collect()
        }
    }

    /// Returns the size of `self`.
    #[inline]
    pub fn size(&self) -> usize { self.data.len() }

    /// Join two sets that contains given keys (Union operation).
    ///
    /// Returns `true` if these keys are belonged to different sets.
    #[inline]
    pub fn union(&mut self, key0: usize, key1: usize) -> bool {
        let k0 = self.get_key(key0);
        let k1 = self.get_key(key1);
        if k0 == k1 { return false; }

        // Temporary replace with dummy to move out the elements of the vector.
        let v0 = mem::replace(&mut self.data[k0], UFNode::Key(usize::MAX)).unwrap();
        let v1 = mem::replace(&mut self.data[k1], UFNode::Key(usize::MAX)).unwrap();

        match UFValue::merge(v0, v1) {
            Merge::Left(val) => {
                self.data[k0] = UFNode::Value(val);
                self.data[k1] = UFNode::Key(k0);
            }
            Merge::Right(val) => {
                self.data[k1] = UFNode::Value(val);
                self.data[k0] = UFNode::Key(k1);
            }
        }

        true
    }

    /// Returns `true` if two keys contained by the same set (find operation).
    #[inline]
    pub fn find(&mut self, key0: usize, key1: usize) -> bool {
         self.get_key(key0) == self.get_key(key1)
    }

    /// Returns the reference to the value of the set that the key belongs to.
    #[inline]
    pub fn get(&mut self, key: usize) -> &T {
        let root_key = self.get_key(key);
        self.data[root_key].get()
    }

    /// Returns the mutable reference to the value of the set that the key belongs to.
    #[inline]
    pub fn get_mut(&mut self, key: usize) -> &mut T {
        let root_key = self.get_key(key);
        self.data[root_key].get_mut()
    }

    fn get_key(&mut self, key: usize) -> usize {
        let root_key = match self.data[key] {
            UFNode::Value(_) => return key,
            UFNode::Key(key) => self.get_key(key)
        };

        self.data[key] = UFNode::Key(root_key);
        root_key
    }
}

impl<A: UFValue> FromIterator<A> for UnionFind<A> {
    #[inline]
    fn from_iter<T: IntoIterator<Item=A>>(iterator: T) -> UnionFind<A> {
        UnionFind { data: iterator.into_iter().map(UFNode::Value).collect() }
    }
}

#[cfg(test)]
mod tests {
    use super::{UnionFind, Size};

    #[test]
    fn test_union_find() {
        let mut uf = UnionFind::new(100);
        assert_eq!(&Size(1), uf.get(0));
        assert_eq!(&Size(1), uf.get(1));
        assert!(!uf.find(0, 1));
        assert!(!uf.find(1, 2));
        assert!(uf.union(0, 1));
        assert!(uf.find(0, 1));
        assert_eq!(&Size(2), uf.get(0));
        assert_eq!(&Size(2), uf.get(1));
        assert_eq!(&Size(1), uf.get(2));
        assert!(!uf.union(0, 1));
        assert_eq!(&Size(2), uf.get(0));
        assert_eq!(&Size(2), uf.get(1));
        assert_eq!(&Size(1), uf.get(2));
        assert!(uf.union(1, 2));
        assert_eq!(&Size(3), uf.get(0));
        assert_eq!(&Size(3), uf.get(1));
        assert_eq!(&Size(3), uf.get(2));
        assert!(uf.find(0, 1));
        assert!(uf.find(2, 1));
    }
}

#[cfg(test)]
mod bench {
    extern crate test;

    use self::test::Bencher;
    use super::{UnionFind, Size};

    fn union_short_trail(uf: &mut UnionFind<Size>) {
        for i in 1 .. uf.size() {
            uf.union(0, i);
        }
    }

    fn union_long_trail(uf: &mut UnionFind<Size>) {
        let mut k = 1;
        while k < uf.size() {
            for i in 0 .. (uf.size() / k / 2) {
                uf.union(2 * k * i, 2 * k * i + k);
            }
            k *= 2;
        }
    }

    fn find_all(uf: &mut UnionFind<Size>) {
        for i in 1 .. uf.size() {
            uf.find(0, i);
        }
    }

    #[bench]
    fn bench_union_short_trail(bencher: &mut Bencher) {
        bencher.iter(|| {
            let mut uf = UnionFind::<Size>::new(1024);
            union_short_trail(&mut uf);
        })
    }

    #[bench]
    fn bench_union_long_trail(bencher: &mut Bencher) {
        bencher.iter(|| {
            let mut uf = UnionFind::<Size>::new(1024);
            union_long_trail(&mut uf);
            find_all(&mut uf);
        })
    }

    #[bench]
    fn bench_find_short_trail(bencher: &mut Bencher) {
        let mut uf = UnionFind::<Size>::new(1024);
        union_short_trail(&mut uf);
        bencher.iter(|| {
            find_all(&mut uf.clone());
        })
    }

    #[bench]
    fn bench_find_long_trail(bencher: &mut Bencher) {
        let mut uf = UnionFind::<Size>::new(1024);
        union_short_trail(&mut uf);
        bencher.iter(|| {
            find_all(&mut uf.clone());
        })
    }
}
<commit_msg>fix warnings<commit_after>//! Struct and methods for union-find operation.

#![warn(bad_style, missing_docs,
        unused, unused_extern_crates, unused_import_braces,
        unused_qualifications, unused_results)]

#![cfg_attr(test, feature(test))]

use std::{mem, usize};
use std::default::Default;
use std::iter::{IntoIterator, FromIterator};

#[derive(Clone, Debug)]
enum UFNode<T> {
    Key(usize),
    Value(T)
}

impl<T> UFNode<T> {
    fn get(&self) -> &T {
        match self {
            &UFNode::Value(ref val) => val,
            &UFNode::Key(_) => panic!()
        }
    }
    fn get_mut(&mut self) -> &mut T {
        match self {
            &mut UFNode::Value(ref mut val) => val,
            &mut UFNode::Key(_) => panic!()
        }
    }
    fn unwrap(self) -> T {
        match self {
            UFNode::Value(val) => val,
            UFNode::Key(_) => panic!()
        }
    }
}

/// The value that can be contained with `UFValue`.
pub trait UFValue: Sized {
    /// Merge two value into one.
    ///
    /// This is used by `UnionFind::union` operation.
    fn merge(lval: Self, rval: Self) -> Merge<Self>;
}

/// Return value of the `UFValue::merege`
#[allow(missing_docs)]
pub enum Merge<T> {
    Left(T), Right(T)
}

/// Reperesents the size of the group.
#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]
pub struct Size(usize);

impl UFValue for Size {
    fn merge(Size(lval): Size, Size(rval): Size) -> Merge<Size> {
        if lval >= rval {
            Merge::Left(Size(lval + rval))
        } else {
            Merge::Right(Size(lval + rval))
        }
    }
}

impl Default for Size {
    fn default() -> Size { Size(1) }
}

/// Struct for union-find operation.
#[derive(Clone, Debug)]
pub struct UnionFind<V = Size> {
    data: Vec<UFNode<V>>
}

impl<T: UFValue = Size> UnionFind<T> {
    /// Creates empty `UnionFind` struct.
    #[inline]
    pub fn new(len: usize) -> UnionFind<T>
        where T: Default
    {
        UnionFind {
            data: (0 .. len)
                .map(|_| UFNode::Value(Default::default()))
                .collect()
        }
    }

    /// Returns the size of `self`.
    #[inline]
    pub fn size(&self) -> usize { self.data.len() }

    /// Join two sets that contains given keys (Union operation).
    ///
    /// Returns `true` if these keys are belonged to different sets.
    #[inline]
    pub fn union(&mut self, key0: usize, key1: usize) -> bool {
        let k0 = self.get_key(key0);
        let k1 = self.get_key(key1);
        if k0 == k1 { return false; }

        // Temporary replace with dummy to move out the elements of the vector.
        let v0 = mem::replace(&mut self.data[k0], UFNode::Key(usize::MAX)).unwrap();
        let v1 = mem::replace(&mut self.data[k1], UFNode::Key(usize::MAX)).unwrap();

        match UFValue::merge(v0, v1) {
            Merge::Left(val) => {
                self.data[k0] = UFNode::Value(val);
                self.data[k1] = UFNode::Key(k0);
            }
            Merge::Right(val) => {
                self.data[k1] = UFNode::Value(val);
                self.data[k0] = UFNode::Key(k1);
            }
        }

        true
    }

    /// Returns `true` if two keys contained by the same set (find operation).
    #[inline]
    pub fn find(&mut self, key0: usize, key1: usize) -> bool {
         self.get_key(key0) == self.get_key(key1)
    }

    /// Returns the reference to the value of the set that the key belongs to.
    #[inline]
    pub fn get(&mut self, key: usize) -> &T {
        let root_key = self.get_key(key);
        self.data[root_key].get()
    }

    /// Returns the mutable reference to the value of the set that the key belongs to.
    #[inline]
    pub fn get_mut(&mut self, key: usize) -> &mut T {
        let root_key = self.get_key(key);
        self.data[root_key].get_mut()
    }

    fn get_key(&mut self, key: usize) -> usize {
        let root_key = match self.data[key] {
            UFNode::Value(_) => return key,
            UFNode::Key(key) => self.get_key(key)
        };

        self.data[key] = UFNode::Key(root_key);
        root_key
    }
}

impl<A: UFValue> FromIterator<A> for UnionFind<A> {
    #[inline]
    fn from_iter<T: IntoIterator<Item=A>>(iterator: T) -> UnionFind<A> {
        UnionFind { data: iterator.into_iter().map(UFNode::Value).collect() }
    }
}

#[cfg(test)]
mod tests {
    use super::{UnionFind, Size};

    #[test]
    fn test_union_find() {
        let mut uf = UnionFind::new(100);
        assert_eq!(&Size(1), uf.get(0));
        assert_eq!(&Size(1), uf.get(1));
        assert!(!uf.find(0, 1));
        assert!(!uf.find(1, 2));
        assert!(uf.union(0, 1));
        assert!(uf.find(0, 1));
        assert_eq!(&Size(2), uf.get(0));
        assert_eq!(&Size(2), uf.get(1));
        assert_eq!(&Size(1), uf.get(2));
        assert!(!uf.union(0, 1));
        assert_eq!(&Size(2), uf.get(0));
        assert_eq!(&Size(2), uf.get(1));
        assert_eq!(&Size(1), uf.get(2));
        assert!(uf.union(1, 2));
        assert_eq!(&Size(3), uf.get(0));
        assert_eq!(&Size(3), uf.get(1));
        assert_eq!(&Size(3), uf.get(2));
        assert!(uf.find(0, 1));
        assert!(uf.find(2, 1));
    }
}

#[cfg(test)]
mod bench {
    extern crate test;

    use self::test::Bencher;
    use super::{UnionFind, Size};

    fn union_short_trail(uf: &mut UnionFind<Size>) {
        for i in 1 .. uf.size() {
            uf.union(0, i);
        }
    }

    fn union_long_trail(uf: &mut UnionFind<Size>) {
        let mut k = 1;
        while k < uf.size() {
            for i in 0 .. (uf.size() / k / 2) {
                uf.union(2 * k * i, 2 * k * i + k);
            }
            k *= 2;
        }
    }

    fn find_all(uf: &mut UnionFind<Size>) {
        for i in 1 .. uf.size() {
            uf.find(0, i);
        }
    }

    #[bench]
    fn bench_union_short_trail(bencher: &mut Bencher) {
        bencher.iter(|| {
            let mut uf = UnionFind::<Size>::new(1024);
            union_short_trail(&mut uf);
        })
    }

    #[bench]
    fn bench_union_long_trail(bencher: &mut Bencher) {
        bencher.iter(|| {
            let mut uf = UnionFind::<Size>::new(1024);
            union_long_trail(&mut uf);
            find_all(&mut uf);
        })
    }

    #[bench]
    fn bench_find_short_trail(bencher: &mut Bencher) {
        let mut uf = UnionFind::<Size>::new(1024);
        union_short_trail(&mut uf);
        bencher.iter(|| {
            find_all(&mut uf.clone());
        })
    }

    #[bench]
    fn bench_find_long_trail(bencher: &mut Bencher) {
        let mut uf = UnionFind::<Size>::new(1024);
        union_short_trail(&mut uf);
        bencher.iter(|| {
            find_all(&mut uf.clone());
        })
    }
}
<|endoftext|>


<commit_before>extern crate cookie as _cookie;
#[macro_use] extern crate error_chain;
#[macro_use] extern crate iron;
extern crate rand;
#[cfg(feature = "redis-backend")] extern crate redis;
#[cfg(feature = "redis-backend")] extern crate r2d2;
#[cfg(feature = "redis-backend")] extern crate r2d2_redis;

use iron::prelude::*;
use iron::middleware::{AroundMiddleware,Handler};
use iron::typemap;

pub mod backends;
pub mod errors;

/// Re-export of the cookie crate
pub mod cookie {
    pub use _cookie::*;
}

/// A simple key-value storage interface that is internally used by `Session`. After request
/// handling the `write` method is called where the session backend has the chance to e.g. set
/// cookies or otherwise modify the response.
pub trait RawSession {
    fn get_raw(&self, key: &str) -> IronResult<Option<String>>;
    fn set_raw(&mut self, key: &str, value: String) -> IronResult<()>;
    fn clear(&mut self) -> IronResult<()>;
    fn write(&self, response: &mut Response) -> IronResult<()>;
}

pub trait SessionBackend: Send + Sync + 'static {
    type S: RawSession;

    /// Parse the session before request handling and return the data in parsed form.
    fn from_request(&self, request: &mut Request) -> Self::S;
}

/// The most important type, the middleware you need to register.
///
/// First parameter is a session backend, which determines how your session data is actually
/// stored. Depending on whether you store your data clientside (signed cookies) or serverside
/// (e.g. Redis and a session key) you're going to have different security properties.
pub struct SessionStorage<B: SessionBackend> {
    backend: B
}

impl<B: SessionBackend> SessionStorage<B> {
    pub fn new(backend: B) -> Self {
        SessionStorage {
            backend: backend
        }
    }
}

/// The high-level interface you use to modify session data. You obtain this object with
/// `request.session()`.
pub struct Session {
    inner: Box<RawSession>
}

/// A typed interface to the string-to-string mapping. Each type represents a key, each instance of
/// a type can be serialized into a value.
pub trait Value: Sized + 'static {
    fn get_key() -> &'static str;
    fn into_raw(self) -> String;
    fn from_raw(value: String) -> Option<Self>;
}

impl Session {
    /// Get a `Value` from the session.
    pub fn get<T: Value + Sized + 'static>(&self) -> IronResult<Option<T>> {
        Ok(try!(self.inner.get_raw(T::get_key())).and_then(T::from_raw))
    }

    /// Set a `Value` in the session.
    pub fn set<T: Value>(&mut self, t: T) -> IronResult<()> {
        self.inner.set_raw(T::get_key(), t.into_raw())
    }

    /// Clear/delete the session
    pub fn clear(&mut self) -> IronResult<()> {
        self.inner.clear()
    }
}

struct SessionKey;
impl typemap::Key for SessionKey { type Value = Session; }

impl<B: SessionBackend> AroundMiddleware for SessionStorage<B> {
    fn around(self, handler: Box<Handler>) -> Box<Handler> {
        Box::new(move |req: &mut Request| -> IronResult<Response> {
            let s = self.backend.from_request(req);
            req.extensions.insert::<SessionKey>(Session {
                inner: Box::new(s)
            });
            let mut res = handler.handle(req);
            let s = req.extensions.remove::<SessionKey>().unwrap();
            match res {
                Ok(ref mut x) => try!(s.inner.write(x)),
                Err(ref mut e) => try!(s.inner.write(&mut e.response))
            };
            res
        })
    }
}

/// The helper trait to obtain your session data from a request.
pub trait SessionRequestExt {
    fn session(&mut self) -> &mut Session;
}

impl<'a, 'b> SessionRequestExt for Request<'a, 'b> {
    fn session(&mut self) -> &mut Session {
        self.extensions.get_mut::<SessionKey>().unwrap()
    }
}

/// A module with some important traits to star-import.
pub mod traits {
    pub use super::{SessionRequestExt};
}

#[cfg(test)]
mod tests {
    #[test]
    fn it_works() {
    }
}
<commit_msg>Only write session if it changed<commit_after>extern crate cookie as _cookie;
#[macro_use] extern crate error_chain;
#[macro_use] extern crate iron;
extern crate rand;
#[cfg(feature = "redis-backend")] extern crate redis;
#[cfg(feature = "redis-backend")] extern crate r2d2;
#[cfg(feature = "redis-backend")] extern crate r2d2_redis;

use iron::prelude::*;
use iron::middleware::{AroundMiddleware,Handler};
use iron::typemap;

pub mod backends;
pub mod errors;

/// Re-export of the cookie crate
pub mod cookie {
    pub use _cookie::*;
}

/// A simple key-value storage interface that is internally used by `Session`. After request
/// handling the `write` method is called where the session backend has the chance to e.g. set
/// cookies or otherwise modify the response.
pub trait RawSession {
    fn get_raw(&self, key: &str) -> IronResult<Option<String>>;
    fn set_raw(&mut self, key: &str, value: String) -> IronResult<()>;
    fn clear(&mut self) -> IronResult<()>;
    fn write(&self, response: &mut Response) -> IronResult<()>;
}

pub trait SessionBackend: Send + Sync + 'static {
    type S: RawSession;

    /// Parse the session before request handling and return the data in parsed form.
    fn from_request(&self, request: &mut Request) -> Self::S;
}

/// The most important type, the middleware you need to register.
///
/// First parameter is a session backend, which determines how your session data is actually
/// stored. Depending on whether you store your data clientside (signed cookies) or serverside
/// (e.g. Redis and a session key) you're going to have different security properties.
pub struct SessionStorage<B: SessionBackend> {
    backend: B
}

impl<B: SessionBackend> SessionStorage<B> {
    pub fn new(backend: B) -> Self {
        SessionStorage {
            backend: backend
        }
    }
}

/// The high-level interface you use to modify session data. You obtain this object with
/// `request.session()`.
pub struct Session {
    inner: Box<RawSession>,
    has_changed: bool,
}

impl Session {
    fn new(s: Box<RawSession>) -> Self {
        Session {
            inner: s,
            has_changed: false
        }
    }
}

/// A typed interface to the string-to-string mapping. Each type represents a key, each instance of
/// a type can be serialized into a value.
pub trait Value: Sized + 'static {
    fn get_key() -> &'static str;
    fn into_raw(self) -> String;
    fn from_raw(value: String) -> Option<Self>;
}

impl Session {
    /// Get a `Value` from the session.
    pub fn get<T: Value + Sized + 'static>(&self) -> IronResult<Option<T>> {
        Ok(try!(self.inner.get_raw(T::get_key())).and_then(T::from_raw))
    }

    /// Set a `Value` in the session.
    pub fn set<T: Value>(&mut self, t: T) -> IronResult<()> {
        // FIXME: Equality check for less unnecessary writes
        self.has_changed = true;
        self.inner.set_raw(T::get_key(), t.into_raw())
    }

    /// Clear/delete the session
    pub fn clear(&mut self) -> IronResult<()> {
        // FIXME: Equality check for less unnecessary writes
        self.has_changed = true;
        self.inner.clear()
    }
}

struct SessionKey;
impl typemap::Key for SessionKey { type Value = Session; }

impl<B: SessionBackend> AroundMiddleware for SessionStorage<B> {
    fn around(self, handler: Box<Handler>) -> Box<Handler> {
        Box::new(move |req: &mut Request| -> IronResult<Response> {
            let s = self.backend.from_request(req);
            req.extensions.insert::<SessionKey>(Session::new(Box::new(s)));
            let mut res = handler.handle(req);
            let s = req.extensions.remove::<SessionKey>().unwrap();
            if s.has_changed {
                match res {
                    Ok(ref mut r) => try!(s.inner.write(r)),
                    Err(ref mut e) => try!(s.inner.write(&mut e.response))
                }
            };
            res
        })
    }
}

/// The helper trait to obtain your session data from a request.
pub trait SessionRequestExt {
    fn session(&mut self) -> &mut Session;
}

impl<'a, 'b> SessionRequestExt for Request<'a, 'b> {
    fn session(&mut self) -> &mut Session {
        self.extensions.get_mut::<SessionKey>().unwrap()
    }
}

/// A module with some important traits to star-import.
pub mod traits {
    pub use super::{SessionRequestExt};
}

#[cfg(test)]
mod tests {
    #[test]
    fn it_works() {
    }
}
<|endoftext|>


<commit_before>//! A library providing a generic connection pool.
#![feature(unsafe_destructor, phase, if_let)]
#![warn(missing_docs)]
#![doc(html_root_url="https://sfackler.github.io/doc")]

#[phase(plugin, link)]
extern crate log;
extern crate serialize;

use std::comm;
use std::collections::RingBuf;
use std::sync::{Arc, Mutex};
use std::fmt;

pub use config::{Config, ConfigError};

mod config;

/// A trait which provides database-specific functionality.
pub trait PoolManager<C, E>: Send+Sync {
    /// Attempts to create a new connection.
    fn connect(&self) -> Result<C, E>;

    /// Determines if the connection is still connected to the database.
    ///
    /// A standard implementation would check if a simple query like `SELECT 1`
    /// succeeds.
    fn is_valid(&self, conn: &mut C) -> Result<(), E>;

    /// *Quickly* determines if the connection is no longer usable.
    ///
    /// This will be called synchronously every time a connection is returned
    /// to the pool, so it should *not* block. If it returns `true`, the
    /// connection will be discarded.
    ///
    /// For example, an implementation might check if the underlying TCP socket
    /// has disconnected. Implementations that do not support this kind of
    /// fast health check may simply return `false`.
    fn has_broken(&self, conn: &mut C) -> bool;
}

/// A trait which handles errors reported by the `PoolManager`.
pub trait ErrorHandler<E>: Send+Sync {
    /// Handles an error.
    fn handle_error(&self, error: E);
}

impl<E> ErrorHandler<E> for Box<ErrorHandler<E>+Sync+Send> {
    fn handle_error(&self, error: E) {
        (**self).handle_error(error)
    }
}

/// An `ErrorHandler` which does nothing.
pub struct NoopErrorHandler;

impl<E> ErrorHandler<E> for NoopErrorHandler {
    fn handle_error(&self, _: E) {}
}

/// An `ErrorHandler` which logs at the error level.
pub struct LoggingErrorHandler;

impl<E> ErrorHandler<E> for LoggingErrorHandler where E: fmt::Show {
    fn handle_error(&self, error: E) {
        error!("{}", error);
    }
}

enum Command<C> {
    AddConnection,
    TestConnection(C),
}

struct PoolInternals<C> {
    conns: RingBuf<C>,
    num_conns: uint,
}

struct InnerPool<C, E, M, H> where C: Send, E: Send, M: PoolManager<C, E>, H: ErrorHandler<E> {
    config: Config,
    manager: M,
    error_handler: H,
    internals: Mutex<PoolInternals<C>>,
}

/// A generic connection pool.
pub struct Pool<C, E, M, H> where C: Send, E: Send, M: PoolManager<C, E>, H: ErrorHandler<E> {
    helper_chan: Mutex<Sender<Command<C>>>,
    inner: Arc<InnerPool<C, E, M, H>>
}

impl<C, E, M, H> Pool<C, E, M, H>
        where C: Send, E: Send, M: PoolManager<C, E>, H: ErrorHandler<E> {
    /// Creates a new connection pool.
    ///
    /// Returns an `Err` value only if `config` is invalid.
    pub fn new(config: Config, manager: M, error_handler: H)
               -> Result<Pool<C, E, M, H>, ConfigError> {
        try!(config.validate());

        let internals = PoolInternals {
            conns: RingBuf::new(),
            num_conns: config.pool_size,
        };

        let inner = Arc::new(InnerPool {
            config: config,
            manager: manager,
            error_handler: error_handler,
            internals: Mutex::new(internals),
        });

        let (sender, receiver) = comm::channel();
        // FIXME :(
        let receiver = Arc::new(Mutex::new(receiver));

        for _ in range(0, config.helper_tasks) {
            let inner = inner.clone();
            let receiver = receiver.clone();
            spawn(proc() helper_task(receiver, inner));
        }

        for _ in range(0, config.pool_size) {
            sender.send(Command::AddConnection);
        }

        Ok(Pool {
            helper_chan: Mutex::new(sender),
            inner: inner,
        })
    }

    /// Retrieves a connection from the pool.
    pub fn get<'a>(&'a self) -> Result<PooledConnection<'a, C, E, M, H>, ()> {
        let mut internals = self.inner.internals.lock();

        loop {
            match internals.conns.pop_front() {
                Some(mut conn) => {
                    drop(internals);

                    if self.inner.config.test_on_check_out {
                        if let Err(e) = self.inner.manager.is_valid(&mut conn) {
                            self.inner.error_handler.handle_error(e);
                            internals = self.inner.internals.lock();
                            internals.num_conns -= 1;
                            continue
                        }
                    }

                    return Ok(PooledConnection {
                        pool: self,
                        conn: Some(conn),
                    })
                }
                None => internals.cond.wait(),
            }
        }
    }

    fn put_back(&self, mut conn: C) {
        // This is specified to be fast, but call it before locking anyways
        let broken = self.inner.manager.has_broken(&mut conn);

        let mut internals = self.inner.internals.lock();
        if broken {
            internals.num_conns -= 1;
        } else {
            internals.conns.push_back(conn);
            internals.cond.signal();
        }
    }
}

fn helper_task<C, E, M, H>(receiver: Arc<Mutex<Receiver<Command<C>>>>,
                           inner: Arc<InnerPool<C, E, M, H>>)
        where C: Send, E: Send, M: PoolManager<C, E>, H: ErrorHandler<E> {
    loop {
        let receiver = receiver.lock();
        let res = receiver.recv_opt();
        drop(receiver);

        match res {
            Ok(Command::AddConnection) => add_connection(&*inner),
            Ok(Command::TestConnection(conn)) => test_connection(&*inner, conn),
            Err(()) => break,
        }
    }
}

fn add_connection<C, E, M, H>(inner: &InnerPool<C, E, M, H>)
        where C: Send, E: Send, M: PoolManager<C, E>, H: ErrorHandler<E> {
    match inner.manager.connect() {
        Ok(conn) => {
            let mut internals = inner.internals.lock();
            internals.conns.push_back(conn);
            internals.num_conns += 1;
            internals.cond.signal();
        }
        Err(err) => inner.error_handler.handle_error(err),
    }
}

fn test_connection<C, E, M, H>(inner: &InnerPool<C, E, M, H>, mut conn: C)
        where C: Send, E: Send, M: PoolManager<C, E>, H: ErrorHandler<E> {
    match inner.manager.is_valid(&mut conn) {
        Ok(()) => inner.internals.lock().conns.push_back(conn),
        Err(e) => {
            inner.error_handler.handle_error(e);
            inner.internals.lock().num_conns -= 1;
        }
    }
}

/// A smart pointer wrapping a connection.
pub struct PooledConnection<'a, C, E, M, H>
        where C: Send, E: Send, M: PoolManager<C, E>, H: ErrorHandler<E> {
    pool: &'a Pool<C, E, M, H>,
    conn: Option<C>,
}

#[unsafe_destructor]
impl<'a, C, E, M, H> Drop for PooledConnection<'a, C, E, M, H>
        where C: Send, E: Send, M: PoolManager<C, E>, H: ErrorHandler<E> {
    fn drop(&mut self) {
        self.pool.put_back(self.conn.take().unwrap());
    }
}

impl<'a, C, E, M, H> Deref<C> for PooledConnection<'a, C, E, M, H>
        where C: Send, E: Send, M: PoolManager<C, E>, H: ErrorHandler<E> {
    fn deref(&self) -> &C {
        self.conn.as_ref().unwrap()
    }
}

impl<'a, C, E, M, H> DerefMut<C> for PooledConnection<'a, C, E, M, H>
        where C: Send, E: Send, M: PoolManager<C, E>, H: ErrorHandler<E> {
    fn deref_mut(&mut self) -> &mut C {
        self.conn.as_mut().unwrap()
    }
}
<commit_msg>Use TaskPool instead of rolling our own<commit_after>//! A library providing a generic connection pool.
#![feature(unsafe_destructor, phase, if_let)]
#![warn(missing_docs)]
#![doc(html_root_url="https://sfackler.github.io/doc")]

#[phase(plugin, link)]
extern crate log;
extern crate serialize;

use std::collections::RingBuf;
use std::sync::{Arc, Mutex, TaskPool};
use std::fmt;

pub use config::{Config, ConfigError};

mod config;

/// A trait which provides database-specific functionality.
pub trait PoolManager<C, E>: Send+Sync {
    /// Attempts to create a new connection.
    fn connect(&self) -> Result<C, E>;

    /// Determines if the connection is still connected to the database.
    ///
    /// A standard implementation would check if a simple query like `SELECT 1`
    /// succeeds.
    fn is_valid(&self, conn: &mut C) -> Result<(), E>;

    /// *Quickly* determines if the connection is no longer usable.
    ///
    /// This will be called synchronously every time a connection is returned
    /// to the pool, so it should *not* block. If it returns `true`, the
    /// connection will be discarded.
    ///
    /// For example, an implementation might check if the underlying TCP socket
    /// has disconnected. Implementations that do not support this kind of
    /// fast health check may simply return `false`.
    fn has_broken(&self, conn: &mut C) -> bool;
}

/// A trait which handles errors reported by the `PoolManager`.
pub trait ErrorHandler<E>: Send+Sync {
    /// Handles an error.
    fn handle_error(&self, error: E);
}

impl<E> ErrorHandler<E> for Box<ErrorHandler<E>+Sync+Send> {
    fn handle_error(&self, error: E) {
        (**self).handle_error(error)
    }
}

/// An `ErrorHandler` which does nothing.
pub struct NoopErrorHandler;

impl<E> ErrorHandler<E> for NoopErrorHandler {
    fn handle_error(&self, _: E) {}
}

/// An `ErrorHandler` which logs at the error level.
pub struct LoggingErrorHandler;

impl<E> ErrorHandler<E> for LoggingErrorHandler where E: fmt::Show {
    fn handle_error(&self, error: E) {
        error!("{}", error);
    }
}

struct PoolInternals<C> {
    conns: RingBuf<C>,
    num_conns: uint,
}

struct InnerPool<C, E, M, H> where C: Send, E: Send, M: PoolManager<C, E>, H: ErrorHandler<E> {
    config: Config,
    manager: M,
    error_handler: H,
    internals: Mutex<PoolInternals<C>>,
}

/// A generic connection pool.
pub struct Pool<C, E, M, H> where C: Send, E: Send, M: PoolManager<C, E>, H: ErrorHandler<E> {
    inner: Arc<InnerPool<C, E, M, H>>,
    task_pool: Mutex<TaskPool>,
}

impl<C, E, M, H> Pool<C, E, M, H>
        where C: Send, E: Send, M: PoolManager<C, E>, H: ErrorHandler<E> {
    /// Creates a new connection pool.
    ///
    /// Returns an `Err` value only if `config` is invalid.
    pub fn new(config: Config, manager: M, error_handler: H)
               -> Result<Pool<C, E, M, H>, ConfigError> {
        try!(config.validate());

        let internals = PoolInternals {
            conns: RingBuf::new(),
            num_conns: config.pool_size,
        };

        let inner = Arc::new(InnerPool {
            config: config,
            manager: manager,
            error_handler: error_handler,
            internals: Mutex::new(internals),
        });

        let pool = Pool {
            inner: inner,
            task_pool: Mutex::new(TaskPool::new(config.helper_tasks)),
        };

        for _ in range(0, config.pool_size) {
            pool.add_connection();
        }

        Ok(pool)
    }

    /// Retrieves a connection from the pool.
    pub fn get<'a>(&'a self) -> Result<PooledConnection<'a, C, E, M, H>, ()> {
        let mut internals = self.inner.internals.lock();

        loop {
            match internals.conns.pop_front() {
                Some(mut conn) => {
                    drop(internals);

                    if self.inner.config.test_on_check_out {
                        if let Err(e) = self.inner.manager.is_valid(&mut conn) {
                            self.inner.error_handler.handle_error(e);
                            internals = self.inner.internals.lock();
                            internals.num_conns -= 1;
                            continue
                        }
                    }

                    return Ok(PooledConnection {
                        pool: self,
                        conn: Some(conn),
                    })
                }
                None => internals.cond.wait(),
            }
        }
    }

    fn put_back(&self, mut conn: C) {
        // This is specified to be fast, but call it before locking anyways
        let broken = self.inner.manager.has_broken(&mut conn);

        let mut internals = self.inner.internals.lock();
        if broken {
            internals.num_conns -= 1;
        } else {
            internals.conns.push_back(conn);
            internals.cond.signal();
        }
    }

    fn add_connection(&self) {
        let inner = self.inner.clone();
        self.task_pool.lock().execute(proc() {
            match inner.manager.connect() {
                Ok(conn) => {
                    let mut internals = inner.internals.lock();
                    internals.conns.push_back(conn);
                    internals.num_conns += 1;
                    internals.cond.signal();
                }
                Err(err) => inner.error_handler.handle_error(err),
            }
        });
    }
}

/// A smart pointer wrapping a connection.
pub struct PooledConnection<'a, C, E, M, H>
        where C: Send, E: Send, M: PoolManager<C, E>, H: ErrorHandler<E> {
    pool: &'a Pool<C, E, M, H>,
    conn: Option<C>,
}

#[unsafe_destructor]
impl<'a, C, E, M, H> Drop for PooledConnection<'a, C, E, M, H>
        where C: Send, E: Send, M: PoolManager<C, E>, H: ErrorHandler<E> {
    fn drop(&mut self) {
        self.pool.put_back(self.conn.take().unwrap());
    }
}

impl<'a, C, E, M, H> Deref<C> for PooledConnection<'a, C, E, M, H>
        where C: Send, E: Send, M: PoolManager<C, E>, H: ErrorHandler<E> {
    fn deref(&self) -> &C {
        self.conn.as_ref().unwrap()
    }
}

impl<'a, C, E, M, H> DerefMut<C> for PooledConnection<'a, C, E, M, H>
        where C: Send, E: Send, M: PoolManager<C, E>, H: ErrorHandler<E> {
    fn deref_mut(&mut self) -> &mut C {
        self.conn.as_mut().unwrap()
    }
}
<|endoftext|>


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

//! This crate contains floating point types that panic if they are set
//! to an illegal value, such as NaN.
//!
//! The name "Noisy Float" comes from
//! the terms "quiet NaN" and "signaling NaN"; "signaling" was too long
//! to put in a struct/crate name, so "noisy" is used instead, being the opposite
//! of "quiet."
//!
//! The standard types defined in `noisy_float::types` follow the principle
//! demonstrated by Rust's handling of integer overflow:
//! a bad arithmetic operation is considered an error,
//! but it is too costly to check everywhere in optimized builds.
//! For each floating point number that is created, a `debug_assert!` invocation is used
//! to check if it is valid or not.
//! This way, there are guarantees when developing code that floating point
//! numbers have valid values,
//! but during a release run there is *no overhead* for using these floating
//! point types compared to using `f32` or `f64` directly.
//!
//! This crate makes use of the floating point and number traits in the
//! popular `num_traits` crate.
//!
//! #Examples
//! An example using the `R64` type, which corresponds to *finite* `f64` values.
//!
//! ```
//! use noisy_float::prelude::*;
//!
//! fn geometric_mean(a: R64, b: R64) -> R64 {
//!     (a * b).sqrt() //used just like regular floating-point numbers
//! }
//!
//! println!("geometric_mean(10.0, 20.0) = {}", geometric_mean(r64(10.0), r64(20.0)));
//! //prints 14.142...
//! ```
//!
//! An example using the `N32` type, which corresponds to *non-NaN* `f32` values.
//! The float types in this crate are able to implement `Eq` and `Ord` properly,
//! since NaN is not allowed.
//!
//! ```
//! use noisy_float::prelude::*;
//!
//! let values = vec![n32(3.0), n32(-1.5), n32(71.3), N32::infinity()];
//! assert!(values.iter().cloned().min() == Some(n32(-1.5)));
//! assert!(values.iter().cloned().max() == Some(N32::infinity()));
//! ```

extern crate num_traits;

pub mod checkers;
pub mod types;
mod float_impl;

/// Prelude for the `noisy_float` crate.
///
/// This includes all of the types defined in the `noisy_float::types` module,
/// as well as re-exports of traits from the `num_traits` crate.
/// It is important to have these re-exports here, because they allow the user
/// to access common floating point methods like `abs()`, `sqrt()`, etc.
pub mod prelude {
    pub use types::*;

    pub use num_traits::Float;
    pub use num_traits::cast::{ToPrimitive, NumCast};
}

use std::marker::PhantomData;
use std::fmt;
use num_traits::Float;

/// Trait for checking whether a floating point number is *valid*.
///
/// The implementation defines its own criteria for what constitutes a *valid* value.
pub trait FloatChecker<F> {
    /// Returns `true` if (and only if) the given floating point number is *valid*
    /// according to this checker's criteria.
    ///
    /// The only hard requirement is that NaN *must* be considered *invalid*
    /// for all implementations of `FloatChecker`.
    fn check(value: F) -> bool;

    /// A function that may panic if the floating point number is *invalid*.
    ///
    /// Should either call `assert!(check(value), ...)` or `debug_assert!(check(value), ...)`.
    fn assert(value: F);
}

/// A floating point number with a restricted set of legal values.
///
/// Typical users will not need to access this struct directly, but
/// can instead use the type aliases found in the module `noisy_float::types`.
/// However, this struct together with a `FloatChecker` implementation can be used
/// to define custom behavior.
///
/// The underlying float type is `F`, usually `f32` or `f64`.
/// Valid values for the float are determined by the float checker `C`.
/// If an invalid value would ever be returned from a method on this type,
/// the method will panic instead, using either `assert!` or `debug_assert!`
/// as defined by the float checker.
/// The exception to this rule is for methods that return an `Option` containing
/// a `NoisyFloat`, in which case the result would be `None` if the value is invalid.
pub struct NoisyFloat<F: Float, C: FloatChecker<F>> {
    value: F,
    checker: PhantomData<C>
}

impl<F: Float, C: FloatChecker<F>> NoisyFloat<F, C> {
    /// Constructs a new `NoisyFloat` with the given value.
    ///
    /// Uses the `FloatChecker` to assert that the value is valid.
    #[inline]
    pub fn new(value: F) -> NoisyFloat<F, C> {
        C::assert(value);
        Self::unchecked_new(value)
    }

    #[inline]
    fn unchecked_new(value: F) -> NoisyFloat<F, C> {
        NoisyFloat {
            value: value,
            checker: PhantomData
        }
    }

    /// Constructs a new `NoisyFloat` with the given value.
    ///
    /// Returns `None` if the value is invalid.
    #[inline]
    pub fn try_new(value: F) -> Option<NoisyFloat<F, C>> {
        if C::check(value) {
            Some(NoisyFloat {
                value: value,
                checker: PhantomData
            })
        } else {
            None
        }
    }

    /// Returns the underlying float value.
    #[inline]
    pub fn raw(self) -> F {
        self.value
    }
}

/// Note: due to complications with Rust's type system, cannot implement `Into` generically like
/// `impl<F: Float, C: FloatChecker<F>> Into<F> for NoisyFloat<F, C>`.
impl<C: FloatChecker<f32>> Into<f32> for NoisyFloat<f32, C> {
    #[inline]
    fn into(self) -> f32 {
        self.value
    }
}

/// Note: due to complications with Rust's type system, cannot implement `Into` generically like
/// `impl<F: Float, C: FloatChecker<F>> Into<F> for NoisyFloat<F, C>`.
impl<C: FloatChecker<f64>> Into<f64> for NoisyFloat<f64, C> {
    #[inline]
    fn into(self) -> f64 {
        self.value
    }
}

impl<F: Float + fmt::Debug, C: FloatChecker<F>> fmt::Debug for NoisyFloat<F, C> {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        fmt::Debug::fmt(&self.value, f)
    }
}

impl<F: Float + fmt::Display, C: FloatChecker<F>> fmt::Display for NoisyFloat<F, C> {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        fmt::Display::fmt(&self.value, f)
    }
}

impl<F: Float + fmt::LowerExp, C: FloatChecker<F>> fmt::LowerExp for NoisyFloat<F, C> {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        fmt::LowerExp::fmt(&self.value, f)
    }
}

impl<F: Float + fmt::UpperExp, C: FloatChecker<F>> fmt::UpperExp for NoisyFloat<F, C> {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        fmt::UpperExp::fmt(&self.value, f)
    }
}

#[cfg(test)]
mod tests {
    use prelude::*;
    use std::f32;
    use std::f64::{self, consts};

    #[test]
    fn smoke_test() {
        assert!(n64(1.0) + n64(2.0) == n64(3.0));
        assert!(n64(3.0) != n64(2.9));
        assert!(r64(1.0) < r64(2.0));
        let mut value = n64(18.0);
        value %= n64(5.0);
        assert!(-value == n64(-3.0));
        assert!(r64(1.0).exp() == r64(consts::E));
        assert!((N64::try_new(1.0).unwrap() / N64::infinity()) == n64(0.0));
        assert!(NumCast::from(f32::INFINITY) == N64::try_new(f64::INFINITY));
        assert!(R64::try_new(f64::NEG_INFINITY) == None);
        assert!(N64::try_new(f64::NAN) == None);
        assert!(R64::try_new(f64::NAN) == None);
    }

    #[test]
    #[should_panic]
    fn n64_nan() {
        n64(0.0) / n64(0.0);
    }

    #[test]
    #[should_panic]
    fn r64_nan() {
        r64(0.0) / r64(0.0);
    }

    #[test]
    #[should_panic]
    fn r64_infinity() {
        r64(1.0) / r64(0.0);
    }
}
<commit_msg>added conversions from primitives to NoisyFloat and reduced prelude<commit_after>// Copyright 2016 Matthew D. Michelotti
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! This crate contains floating point types that panic if they are set
//! to an illegal value, such as NaN.
//!
//! The name "Noisy Float" comes from
//! the terms "quiet NaN" and "signaling NaN"; "signaling" was too long
//! to put in a struct/crate name, so "noisy" is used instead, being the opposite
//! of "quiet."
//!
//! The standard types defined in `noisy_float::types` follow the principle
//! demonstrated by Rust's handling of integer overflow:
//! a bad arithmetic operation is considered an error,
//! but it is too costly to check everywhere in optimized builds.
//! For each floating point number that is created, a `debug_assert!` invocation is used
//! to check if it is valid or not.
//! This way, there are guarantees when developing code that floating point
//! numbers have valid values,
//! but during a release run there is *no overhead* for using these floating
//! point types compared to using `f32` or `f64` directly.
//!
//! This crate makes use of the floating point and number traits in the
//! popular `num_traits` crate.
//!
//! #Examples
//! An example using the `R64` type, which corresponds to *finite* `f64` values.
//!
//! ```
//! use noisy_float::prelude::*;
//!
//! fn geometric_mean(a: R64, b: R64) -> R64 {
//!     (a * b).sqrt() //used just like regular floating-point numbers
//! }
//!
//! println!("geometric_mean(10.0, 20.0) = {}", geometric_mean(r64(10.0), r64(20.0)));
//! //prints 14.142...
//! ```
//!
//! An example using the `N32` type, which corresponds to *non-NaN* `f32` values.
//! The float types in this crate are able to implement `Eq` and `Ord` properly,
//! since NaN is not allowed.
//!
//! ```
//! use noisy_float::prelude::*;
//!
//! let values = vec![n32(3.0), n32(-1.5), n32(71.3), N32::infinity()];
//! assert!(values.iter().cloned().min() == Some(n32(-1.5)));
//! assert!(values.iter().cloned().max() == Some(N32::infinity()));
//! ```

extern crate num_traits;

pub mod checkers;
pub mod types;
mod float_impl;

/// Prelude for the `noisy_float` crate.
///
/// This includes all of the types defined in the `noisy_float::types` module,
/// as well as a re-export of the `Float` trait from the `num_traits` crate.
/// It is important to have this re-export here, because it allows the user
/// to access common floating point methods like `abs()`, `sqrt()`, etc.
pub mod prelude {
    pub use types::*;

    pub use num_traits::Float;
}

use std::marker::PhantomData;
use std::fmt;
use num_traits::Float;

/// Trait for checking whether a floating point number is *valid*.
///
/// The implementation defines its own criteria for what constitutes a *valid* value.
pub trait FloatChecker<F> {
    /// Returns `true` if (and only if) the given floating point number is *valid*
    /// according to this checker's criteria.
    ///
    /// The only hard requirement is that NaN *must* be considered *invalid*
    /// for all implementations of `FloatChecker`.
    fn check(value: F) -> bool;

    /// A function that may panic if the floating point number is *invalid*.
    ///
    /// Should either call `assert!(check(value), ...)` or `debug_assert!(check(value), ...)`.
    fn assert(value: F);
}

/// A floating point number with a restricted set of legal values.
///
/// Typical users will not need to access this struct directly, but
/// can instead use the type aliases found in the module `noisy_float::types`.
/// However, this struct together with a `FloatChecker` implementation can be used
/// to define custom behavior.
///
/// The underlying float type is `F`, usually `f32` or `f64`.
/// Valid values for the float are determined by the float checker `C`.
/// If an invalid value would ever be returned from a method on this type,
/// the method will panic instead, using either `assert!` or `debug_assert!`
/// as defined by the float checker.
/// The exception to this rule is for methods that return an `Option` containing
/// a `NoisyFloat`, in which case the result would be `None` if the value is invalid.
pub struct NoisyFloat<F: Float, C: FloatChecker<F>> {
    value: F,
    checker: PhantomData<C>
}

impl<F: Float, C: FloatChecker<F>> NoisyFloat<F, C> {
    /// Constructs a `NoisyFloat` with the given value.
    ///
    /// Uses the `FloatChecker` to assert that the value is valid.
    #[inline]
    pub fn new(value: F) -> Self {
        C::assert(value);
        Self::unchecked_new(value)
    }

    #[inline]
    fn unchecked_new(value: F) -> Self {
        NoisyFloat {
            value: value,
            checker: PhantomData
        }
    }

    /// Tries to construct a `NoisyFloat` with the given value.
    ///
    /// Returns `None` if the value is invalid.
    #[inline]
    pub fn try_new(value: F) -> Option<Self> {
        if C::check(value) {
            Some(NoisyFloat {
                value: value,
                checker: PhantomData
            })
        } else {
            None
        }
    }

    /// Constructs a `NoisyFloat` with the given `f32` value.
    ///
    /// May panic not only by the `FloatChecker` but also
    /// by unwrapping the result of a `NumCast` invocation for type `F`,
    /// although the later should not occur in normal situations.
    #[inline]
    pub fn from_f32(value: f32) -> Self {
        Self::new(F::from(value).unwrap())
    }

    /// Constructs a `NoisyFloat` with the given `f64` value.
    ///
    /// May panic not only by the `FloatChecker` but also
    /// by unwrapping the result of a `NumCast` invocation for type `F`,
    /// although the later should not occur in normal situations.
    #[inline]
    pub fn from_f64(value: f64) -> Self {
        Self::new(F::from(value).unwrap())
    }

    /// Returns the underlying float value.
    #[inline]
    pub fn raw(self) -> F {
        self.value
    }
}

/// Note: due to complications with Rust's type system, cannot implement `Into` generically like
/// `impl<F: Float, C: FloatChecker<F>> Into<F> for NoisyFloat<F, C>`.
impl<C: FloatChecker<f32>> Into<f32> for NoisyFloat<f32, C> {
    #[inline]
    fn into(self) -> f32 {
        self.value
    }
}

/// Note: due to complications with Rust's type system, cannot implement `Into` generically like
/// `impl<F: Float, C: FloatChecker<F>> Into<F> for NoisyFloat<F, C>`.
impl<C: FloatChecker<f64>> Into<f64> for NoisyFloat<f64, C> {
    #[inline]
    fn into(self) -> f64 {
        self.value
    }
}

impl<F: Float + fmt::Debug, C: FloatChecker<F>> fmt::Debug for NoisyFloat<F, C> {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        fmt::Debug::fmt(&self.value, f)
    }
}

impl<F: Float + fmt::Display, C: FloatChecker<F>> fmt::Display for NoisyFloat<F, C> {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        fmt::Display::fmt(&self.value, f)
    }
}

impl<F: Float + fmt::LowerExp, C: FloatChecker<F>> fmt::LowerExp for NoisyFloat<F, C> {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        fmt::LowerExp::fmt(&self.value, f)
    }
}

impl<F: Float + fmt::UpperExp, C: FloatChecker<F>> fmt::UpperExp for NoisyFloat<F, C> {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        fmt::UpperExp::fmt(&self.value, f)
    }
}

#[cfg(test)]
mod tests {
    use prelude::*;
    use std::f32;
    use std::f64::{self, consts};

    #[test]
    fn smoke_test() {
        assert!(n64(1.0) + n64(2.0) == n64(3.0));
        assert!(n64(3.0) != n64(2.9));
        assert!(r64(1.0) < r64(2.0));
        let mut value = n64(18.0);
        value %= n64(5.0);
        assert!(-value == n64(-3.0));
        assert!(r64(1.0).exp() == r64(consts::E));
        assert!((N64::try_new(1.0).unwrap() / N64::infinity()) == n64(0.0));
        assert!(N64::from_f32(f32::INFINITY) == N64::from_f64(f64::INFINITY));
        assert!(R64::try_new(f64::NEG_INFINITY) == None);
        assert!(N64::try_new(f64::NAN) == None);
        assert!(R64::try_new(f64::NAN) == None);
    }

    #[test]
    #[should_panic]
    fn n64_nan() {
        n64(0.0) / n64(0.0);
    }

    #[test]
    #[should_panic]
    fn r64_nan() {
        r64(0.0) / r64(0.0);
    }

    #[test]
    #[should_panic]
    fn r64_infinity() {
        r64(1.0) / r64(0.0);
    }
}
<|endoftext|>


<commit_before>#![crate_id = "graphics"]
#![deny(missing_doc)]

//! Attempt of creating a cheap drawing context.

use std::num::Float;

pub type Matrix2d = [f64, ..6];
pub type Color = [f64, ..4];
pub type Rectangle = [f64, ..4];

/// Implemented by all graphics back-ends.
/// This trait uses default methods to simplify implementation.
pub trait BackEnd {
    /// Returns true if feature is supported.
    #[inline(always)]
    fn supports_tri_list_xy_rgba_f64(&self) -> bool { false }

    /// Renders list of 2d triangles with color assigned per vertex.
    fn tri_list_xy_rgba_f64(&mut self, _vertices: &[f64], _colors: &[f64]) {}

    /// Returns true if feature is supported.
    #[inline(always)]
    fn supports_tri_list_xy_rgba_f32(&self) -> bool { false }

    /// Renders list of 2d triangles with color assigned per vertex.
    fn tri_list_xy_rgba_f32(&mut self, _vertices: &[f32], _colors: &[f32]) {}
}

/// Implemented by contexts that can transform.
pub trait Transform2d<'a> {
    /// Translate x and y.
    fn trans(&'a self, x: f64, y: f64) -> Self;
    /// Rotate degrees.
    fn rot_deg(&'a self, angle: f64) -> Self;
    /// Rotate radians.
    fn rot_rad(&'a self, angle: f64) -> Self;
    /// Scale.
    fn scale(&'a self, sx: f64, sy: f64) -> Self;
    /// Shear.
    fn shear(&'a self, sx: f64, sy: f64) -> Self;
}

/// A structure that might contain a value or a borrowed value.
/// This is to used as building block to create data structure
/// that is partially based on an existing structure.
pub enum Field<'a, T> {
    /// Contains a value.
    Value(T),
    /// Contains a borrowed pointer.
    Borrowed(&'a T),
}

impl<'a, T> Field<'a, T> {
    /// Gets a read only value.
    #[inline(always)]
    pub fn get(&'a self) -> &'a T {
        match *self {
            Value(ref val) => val,
            Borrowed(rval) => rval,
        }
    }
}

/// Multiplies two matrices.
#[inline(always)]
pub fn multiply(m: &[f64, ..6], b: &[f64, ..6]) -> [f64, ..6] {
    [m[0]*b[0]+m[1]*b[3]+m[2]*0.0,  m[0]*b[1]+m[1]*b[4]+m[2]*0.0,  m[0]*b[2]+m[1]*b[5]+m[2]*1.0,
     m[3]*b[0]+m[4]*b[3]+m[5]*0.0,  m[3]*b[1]+m[4]*b[4]+m[5]*0.0,  m[3]*b[2]+m[4]*b[5]+m[5]*1.0]
}

/// Drawing 2d context.
pub struct Context<'a> {
    base: Field<'a, Matrix2d>,
    transform: Field<'a, Matrix2d>,
    color: Field<'a, Color>,
}

impl<'a> Transform2d<'a> for Context<'a> {
    #[inline(always)]
    fn trans(&'a self, x: f64, y: f64) -> Context<'a> {
        Context {
            base: Borrowed(self.base.get()),
            transform: Value({
                let trans: [f64, ..6] = [1.0, 0.0, x,
                                         0.0, 1.0, y];
                 multiply(&trans, self.transform.get())
            }),
            color: Borrowed(self.color.get()),
        }
    }

    #[inline(always)]
    fn rot_deg(&'a self, angle: f64) -> Context<'a> {
        let pi: f64 = Float::pi();
        self.rot_rad(angle * pi / 180.0)
    }

    #[inline(always)]
    fn rot_rad(&'a self, angle: f64) -> Context<'a> {
        Context {
            base: Borrowed(self.base.get()),
            transform: Value({
                let c = angle.cos();
                let s = angle.sin();
                let rot: [f64, ..6] = [c, s, 0.0,
                                      -s, c, 0.0];
                multiply(&rot, self.transform.get())
            }),
            color: Borrowed(self.color.get()),
        }
    }

    #[inline(always)]
    fn scale(&'a self, sx: f64, sy: f64) -> Context<'a> {
        Context {
            base: Borrowed(self.base.get()),
            transform: Value({
                let scale: [f64, ..6] = [sx, 0.0, 0.0,
                                         0.0, sy, 0.0];
                multiply(&scale, self.transform.get())
            }),
            color: Borrowed(self.color.get()),
        }
    }

    #[inline(always)]
    fn shear(&'a self, sx: f64, sy: f64) -> Context<'a> {
        Context {
            base: Borrowed(self.base.get()),
            transform: Value({
                let shear: [f64, ..6] = [1.0, sx, 0.0,
                                         sy, 1.0, 0.0];
            multiply(&shear, self.transform.get())
            }),
            color: Borrowed(self.color.get()),
        }
    }
}

impl<'a> Context<'a> {
    /// Creates a new drawing context.
    pub fn new() -> Context {
        Context {
            base:  Value([1.0, 0.0, 0.0,
                          0.0, 1.0, 0.0]),
            transform: Value([1.0, 0.0, 0.0,
                          0.0, 1.0, 0.0]),
            color: Value([0.0, 0.0, 0.0, 1.0]),
        }
    }

}

#[test]
fn test_context() {
    let c = Context::new();
    {
        let d = c.trans(20.0, 40.0);
        let d = d.trans(10.0, 10.0);
        assert_eq!(d.transform.get()[2], 30.0);
        assert_eq!(d.transform.get()[5], 50.0);
    }
    assert_eq!(c.transform.get()[2], 0.0);
    assert_eq!(c.transform.get()[5], 0.0);

    let c = c.rot_deg(90.0);
    assert!((c.transform.get()[0] - 0.0).abs() < 0.00001);
    assert!((c.transform.get()[1] - 1.0).abs() < 0.00001);
}

#[test]
fn test_scale() {
    let c = Context::new();
    let c = c.scale(2.0, 3.0);
    assert!((c.transform.get()[0] - 2.0).abs() < 0.00001);
    assert!((c.transform.get()[4] - 3.0).abs() < 0.00001);
}

/// A context with color information.
pub struct ColorContext<'a> {
    base: Field<'a, Matrix2d>,
    transform: Field<'a, Matrix2d>,
    color: Field<'a, Color>,
}

/// A rectangle context.
pub struct RectangleContext<'a> {
    base: Field<'a, Matrix2d>,
    transform: Field<'a, Matrix2d>,
    rect: Field<'a, Rectangle>,
}

/// A rectangle color context.
pub struct RectangleColorContext<'a> {
    base: Field<'a, Matrix2d>,
    transform: Field<'a, Matrix2d>,
    rect: Field<'a, Rectangle>,
    color: Field<'a, Color>,
}

<commit_msg>Deleted colour from Context<commit_after>#![crate_id = "graphics"]
#![deny(missing_doc)]

//! Attempt of creating a cheap drawing context.

use std::num::Float;

pub type Matrix2d = [f64, ..6];
pub type Color = [f64, ..4];
pub type Rectangle = [f64, ..4];

/// Implemented by all graphics back-ends.
/// This trait uses default methods to simplify implementation.
pub trait BackEnd {
    /// Returns true if feature is supported.
    #[inline(always)]
    fn supports_tri_list_xy_rgba_f64(&self) -> bool { false }

    /// Renders list of 2d triangles with color assigned per vertex.
    fn tri_list_xy_rgba_f64(&mut self, _vertices: &[f64], _colors: &[f64]) {}

    /// Returns true if feature is supported.
    #[inline(always)]
    fn supports_tri_list_xy_rgba_f32(&self) -> bool { false }

    /// Renders list of 2d triangles with color assigned per vertex.
    fn tri_list_xy_rgba_f32(&mut self, _vertices: &[f32], _colors: &[f32]) {}
}

/// Implemented by contexts that can transform.
pub trait Transform2d<'a> {
    /// Translate x and y.
    fn trans(&'a self, x: f64, y: f64) -> Self;
    /// Rotate degrees.
    fn rot_deg(&'a self, angle: f64) -> Self;
    /// Rotate radians.
    fn rot_rad(&'a self, angle: f64) -> Self;
    /// Scale.
    fn scale(&'a self, sx: f64, sy: f64) -> Self;
    /// Shear.
    fn shear(&'a self, sx: f64, sy: f64) -> Self;
}

/// A structure that might contain a value or a borrowed value.
/// This is to used as building block to create data structure
/// that is partially based on an existing structure.
pub enum Field<'a, T> {
    /// Contains a value.
    Value(T),
    /// Contains a borrowed pointer.
    Borrowed(&'a T),
}

impl<'a, T> Field<'a, T> {
    /// Gets a read only value.
    #[inline(always)]
    pub fn get(&'a self) -> &'a T {
        match *self {
            Value(ref val) => val,
            Borrowed(rval) => rval,
        }
    }
}

/// Multiplies two matrices.
#[inline(always)]
pub fn multiply(m: &[f64, ..6], b: &[f64, ..6]) -> [f64, ..6] {
    [m[0]*b[0]+m[1]*b[3]+m[2]*0.0,  m[0]*b[1]+m[1]*b[4]+m[2]*0.0,  m[0]*b[2]+m[1]*b[5]+m[2]*1.0,
     m[3]*b[0]+m[4]*b[3]+m[5]*0.0,  m[3]*b[1]+m[4]*b[4]+m[5]*0.0,  m[3]*b[2]+m[4]*b[5]+m[5]*1.0]
}

/// Drawing 2d context.
pub struct Context<'a> {
    base: Field<'a, Matrix2d>,
    transform: Field<'a, Matrix2d>,
}

impl<'a> Transform2d<'a> for Context<'a> {
    #[inline(always)]
    fn trans(&'a self, x: f64, y: f64) -> Context<'a> {
        Context {
            base: Borrowed(self.base.get()),
            transform: Value({
                let trans: [f64, ..6] = [1.0, 0.0, x,
                                         0.0, 1.0, y];
                 multiply(&trans, self.transform.get())
            }),
        }
    }

    #[inline(always)]
    fn rot_deg(&'a self, angle: f64) -> Context<'a> {
        let pi: f64 = Float::pi();
        self.rot_rad(angle * pi / 180.0)
    }

    #[inline(always)]
    fn rot_rad(&'a self, angle: f64) -> Context<'a> {
        Context {
            base: Borrowed(self.base.get()),
            transform: Value({
                let c = angle.cos();
                let s = angle.sin();
                let rot: [f64, ..6] = [c, s, 0.0,
                                      -s, c, 0.0];
                multiply(&rot, self.transform.get())
            }),
        }
    }

    #[inline(always)]
    fn scale(&'a self, sx: f64, sy: f64) -> Context<'a> {
        Context {
            base: Borrowed(self.base.get()),
            transform: Value({
                let scale: [f64, ..6] = [sx, 0.0, 0.0,
                                         0.0, sy, 0.0];
                multiply(&scale, self.transform.get())
            }),
        }
    }

    #[inline(always)]
    fn shear(&'a self, sx: f64, sy: f64) -> Context<'a> {
        Context {
            base: Borrowed(self.base.get()),
            transform: Value({
                let shear: [f64, ..6] = [1.0, sx, 0.0,
                                         sy, 1.0, 0.0];
            multiply(&shear, self.transform.get())
            }),
        }
    }
}

impl<'a> Context<'a> {
    /// Creates a new drawing context.
    pub fn new() -> Context {
        Context {
            base:  Value([1.0, 0.0, 0.0,
                          0.0, 1.0, 0.0]),
            transform: Value([1.0, 0.0, 0.0,
                          0.0, 1.0, 0.0]),
        }
    }

}

#[test]
fn test_context() {
    let c = Context::new();
    {
        let d = c.trans(20.0, 40.0);
        let d = d.trans(10.0, 10.0);
        assert_eq!(d.transform.get()[2], 30.0);
        assert_eq!(d.transform.get()[5], 50.0);
    }
    assert_eq!(c.transform.get()[2], 0.0);
    assert_eq!(c.transform.get()[5], 0.0);

    let c = c.rot_deg(90.0);
    assert!((c.transform.get()[0] - 0.0).abs() < 0.00001);
    assert!((c.transform.get()[1] - 1.0).abs() < 0.00001);
}

#[test]
fn test_scale() {
    let c = Context::new();
    let c = c.scale(2.0, 3.0);
    assert!((c.transform.get()[0] - 2.0).abs() < 0.00001);
    assert!((c.transform.get()[4] - 3.0).abs() < 0.00001);
}

/// A context with color information.
pub struct ColorContext<'a> {
    base: Field<'a, Matrix2d>,
    transform: Field<'a, Matrix2d>,
    color: Field<'a, Color>,
}

/// A rectangle context.
pub struct RectangleContext<'a> {
    base: Field<'a, Matrix2d>,
    transform: Field<'a, Matrix2d>,
    rect: Field<'a, Rectangle>,
}

/// A rectangle color context.
pub struct RectangleColorContext<'a> {
    base: Field<'a, Matrix2d>,
    transform: Field<'a, Matrix2d>,
    rect: Field<'a, Rectangle>,
    color: Field<'a, Color>,
}

<|endoftext|>


<commit_before>/// Implementation of Stochastic universal sampling
/// https://en.wikipedia.org/wiki/Stochastic_universal_sampling
/// Runtime: O(n)
/// Memory Usage: O(n) or O(1) (in place variant)

extern crate rand;

use self::rand::{thread_rng, Rng};

pub struct RandomChoice;

impl RandomChoice {
    /// Chooses n samples by their weights. The greater their weights the more likely they get chosen.
    ///
    /// @invariant sum of weights must not overflow.
    /// @param samples The to be selected samples
    /// @param weights Weights that get chosen by their weight/probability. One weight can be greater 1.
    /// @param n Number of randomly chosen samples by weight.
    /// @return randomly selected samples by their weights
    pub fn random_choice_f64<'a, T>(samples: &'a [T], weights: &[f64], n: usize) -> Vec<&'a T> {
        if weights.len() == 0 || n == 0 {
            return Vec::new();
        }

        let sum: f64 = weights.iter().fold(0.0, |acc, &i| acc + i);
        let spoke_gap: f64 = sum / n as f64;

        // next_f64() ∈ [0.0, 1.0)
        let spin = thread_rng().next_f64() * spoke_gap;

        let mut i: usize = 0;
        let mut accumulated_weights = weights[0];
        let mut choices: Vec<&T> = Vec::with_capacity(n);
        let mut current_spoke: f64 = spin;

        for _ in 0..n {
            while accumulated_weights < current_spoke {
                i += 1;
                accumulated_weights += weights[i];
            }
            choices.push(&samples[i]);
            current_spoke += spoke_gap;
        }

        choices
    }

    /// Chooses n samples by their weights. The greater their weights the more likely they get chosen.
    /// The result gets saved directly in the samples argument.
    /// @invariant sum of weights must not overflow.
    /// @param samples The to be selected samples
    /// @param weights Weights that get chosen by their weight/probability. One weight can be greater 1.
    pub fn random_choice_in_place_f64<T: Clone>(samples: &mut [T], weights: &[f64]) {
        if weights.len() < 2 {
            return;
        }

        let sum: f64 = weights.iter().fold(0.0, |acc, &i| acc + i);
        let n: usize = weights.len();
        let spoke_gap: f64 = sum / n as f64;

        // next_f64() ∈ [0.0, 1.0)
        let spin = thread_rng().next_f64() * spoke_gap;

        let mut j: usize = 0;
        let mut accumulated_weights = weights[0];
        let mut current_spoke: f64 = spin;

        for i in 0..n {
            while accumulated_weights < current_spoke {
                j += 1;
                accumulated_weights += weights[j];
            }
            samples[i] = samples[j].clone();
            current_spoke += spoke_gap;
        }
    }

    pub fn random_choice_f32<'a, T>(samples: &'a [T], weights: &[f32], n: usize) -> Vec<&'a T> {
        if weights.len() == 0 || n == 0 {
            return Vec::new();
        }

        let sum: f32 = weights.iter().fold(0.0, |acc, &i| acc + i);
        let spoke_gap: f32 = sum / n as f32;

        // next_f32() ∈ [0.0, 1.0)
        let spin = thread_rng().next_f32() * spoke_gap;

        let mut i: usize = 0;
        let mut accumulated_weights = weights[0];
        let mut choices: Vec<&T> = Vec::with_capacity(n);
        let mut current_spoke: f32 = spin;

        for _ in 0..n {
            while accumulated_weights < current_spoke {
                i += 1;
                accumulated_weights += weights[i];
            }
            choices.push(&samples[i]);
            current_spoke += spoke_gap;
        }

        choices
    }

    pub fn random_choice_in_place_f32<T: Clone>(samples: &mut [T], weights: &[f32]) {
        if weights.len() < 2 {
            return;
        }

        let sum: f32 = weights.iter().fold(0.0, |acc, &i| acc + i);
        let n: usize = weights.len();
        let spoke_gap: f32 = sum / n as f32;

        // next_f32() ∈ [0.0, 1.0)
        let spin = thread_rng().next_f32() * spoke_gap;

        let mut j: usize = 0;
        let mut accumulated_weights = weights[0];
        let mut current_spoke: f32 = spin;

        for i in 0..n {
            while accumulated_weights < current_spoke {
                j += 1;
                accumulated_weights += weights[j];
            }
            samples[i] = samples[j].clone();
            current_spoke += spoke_gap;
        }
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn test_random_choice_64() {
        let test_vec: Vec<f64> = vec![3.1, 1.2, 1.3, 2.0];

        let choices = super::RandomChoice::random_choice_f64(&test_vec, &test_vec, 4 as usize);

        for choice in choices {
            print!("{}, ", choice);
        }
    }

    #[test]
    fn test_random_choice_in_place_64() {
        let mut samples = vec!["hi", "this", "is", "a", "test!"];
        let weights: Vec<f64> = vec![1.0, 1.0, 1.0, 1.0, 1.0];

        super::RandomChoice::random_choice_in_place_f64(&mut samples, &weights);

        for sample in samples {
            print!("{}, ", sample);
        }
    }

}<commit_msg>Benches added<commit_after>#![cfg_attr(test, feature(test))]

/// Implementation of Stochastic universal sampling
/// https://en.wikipedia.org/wiki/Stochastic_universal_sampling
/// Runtime: O(n)
/// Memory Usage: O(n) or O(1) (in place variant)

extern crate rand;

use self::rand::{thread_rng, Rng};

pub struct RandomChoice;

impl RandomChoice {
    /// Chooses n samples by their weights. The greater their weights the more likely they get chosen.
    ///
    /// @invariant sum of weights must not overflow.
    /// @param samples The to be selected samples
    /// @param weights Weights that get chosen by their weight/probability. One weight can be greater 1.
    /// @param n Number of randomly chosen samples by weight.
    /// @return randomly selected samples by their weights
    pub fn random_choice_f64<'a, T>(samples: &'a [T], weights: &[f64], n: usize) -> Vec<&'a T> {
        if weights.len() == 0 || n == 0 {
            return Vec::new();
        }

        let sum: f64 = weights.iter().fold(0.0, |acc, &i| acc + i);
        let spoke_gap: f64 = sum / n as f64;

        // next_f64() ∈ [0.0, 1.0)
        let spin = thread_rng().next_f64() * spoke_gap;

        let mut i: usize = 0;
        let mut accumulated_weights = weights[0];
        let mut choices: Vec<&T> = Vec::with_capacity(n);
        let mut current_spoke: f64 = spin;

        for _ in 0..n {
            while accumulated_weights < current_spoke {
                i += 1;
                accumulated_weights += weights[i];
            }
            choices.push(&samples[i]);
            current_spoke += spoke_gap;
        }

        choices
    }

    /// Chooses n samples by their weights. The greater their weights the more likely they get chosen.
    /// The result gets saved directly in the samples argument.
    /// @invariant sum of weights must not overflow.
    /// @param samples The to be selected samples
    /// @param weights Weights that get chosen by their weight/probability. One weight can be greater 1.
    pub fn random_choice_in_place_f64<T: Clone>(samples: &mut [T], weights: &[f64]) {
        if weights.len() < 2 {
            return;
        }

        let sum: f64 = weights.iter().fold(0.0, |acc, &i| acc + i);
        let n: usize = weights.len();
        let spoke_gap: f64 = sum / n as f64;

        // next_f64() ∈ [0.0, 1.0)
        let spin = thread_rng().next_f64() * spoke_gap;

        let mut j: usize = 0;
        let mut accumulated_weights = weights[0];
        let mut current_spoke: f64 = spin;

        for i in 0..n {
            while accumulated_weights < current_spoke {
                j += 1;
                accumulated_weights += weights[j];
            }
            samples[i] = samples[j].clone();
            current_spoke += spoke_gap;
        }
    }

    pub fn random_choice_f32<'a, T>(samples: &'a [T], weights: &[f32], n: usize) -> Vec<&'a T> {
        if weights.len() == 0 || n == 0 {
            return Vec::new();
        }

        let sum: f32 = weights.iter().fold(0.0, |acc, &i| acc + i);
        let spoke_gap: f32 = sum / n as f32;

        // next_f32() ∈ [0.0, 1.0)
        let spin = thread_rng().next_f32() * spoke_gap;

        let mut i: usize = 0;
        let mut accumulated_weights = weights[0];
        let mut choices: Vec<&T> = Vec::with_capacity(n);
        let mut current_spoke: f32 = spin;

        for _ in 0..n {
            while accumulated_weights < current_spoke {
                i += 1;
                accumulated_weights += weights[i];
            }
            choices.push(&samples[i]);
            current_spoke += spoke_gap;
        }

        choices
    }

    pub fn random_choice_in_place_f32<T: Clone>(samples: &mut [T], weights: &[f32]) {
        if weights.len() < 2 {
            return;
        }

        let sum: f32 = weights.iter().fold(0.0, |acc, &i| acc + i);
        let n: usize = weights.len();
        let spoke_gap: f32 = sum / n as f32;

        // next_f32() ∈ [0.0, 1.0)
        let spin = thread_rng().next_f32() * spoke_gap;

        let mut j: usize = 0;
        let mut accumulated_weights = weights[0];
        let mut current_spoke: f32 = spin;

        for i in 0..n {
            while accumulated_weights < current_spoke {
                j += 1;
                accumulated_weights += weights[j];
            }
            samples[i] = samples[j].clone();
            current_spoke += spoke_gap;
        }
    }
}

#[cfg(test)]
mod benches {

    extern crate test;
    use self::test::Bencher;

    #[bench]
    fn bench_random_choice_64(b: &mut Bencher) {
        let capacity: usize = 500;
        let mut samples: Vec<f64> = Vec::with_capacity(capacity);
        let mut weights: Vec<f64> = Vec::with_capacity(capacity);

        for i in 0..capacity {
            samples.push((i + 1usize) as f64);
            weights.push((i + 1usize) as f64);
        }
        b.iter(|| {
            super::RandomChoice::random_choice_f64(&samples, &weights, 1200 as usize);
        });
    }

    #[bench]
    fn bench_random_choice_in_place_64(b: &mut Bencher) {
        let capacity: usize = 500;
        let mut samples: Vec<f64> = Vec::with_capacity(capacity);
        let mut weights: Vec<f64> = Vec::with_capacity(capacity);

        for i in 0..capacity {
            samples.push((i + 1usize) as f64);
            weights.push((i + 1usize) as f64);
        }
        b.iter(|| {
            super::RandomChoice::random_choice_in_place_f64(&mut samples, &weights);
        });
    }

    #[bench]
    fn bench_random_choice_32(b: &mut Bencher) {
        let capacity: usize = 500;
        let mut samples: Vec<f32> = Vec::with_capacity(capacity);
        let mut weights: Vec<f32> = Vec::with_capacity(capacity);

        for i in 0..capacity {
            samples.push((i + 1usize) as f32);
            weights.push((i + 1usize) as f32);
        }
        b.iter(|| {
            super::RandomChoice::random_choice_f32(&samples, &weights, 1200 as usize);
        });
    }

    #[bench]
    fn bench_random_choice_in_place_32(b: &mut Bencher) {
        let capacity: usize = 500;
        let mut samples: Vec<f32> = Vec::with_capacity(capacity);
        let mut weights: Vec<f32> = Vec::with_capacity(capacity);

        for i in 0..capacity {
            samples.push((i + 1usize) as f32);
            weights.push((i + 1usize) as f32);
        }
        b.iter(|| {
            super::RandomChoice::random_choice_in_place_f32(&mut samples, &weights);
        });
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    #[test]
    fn test_random_choice_64() {
        let capacity: usize = 1000;
        let mut samples: Vec<usize> = Vec::with_capacity(capacity);
        let mut weights: Vec<f64> = Vec::with_capacity(capacity);

        for i in 0..capacity {
            samples.push(i + 1);
            weights.push((i + 1usize) as f64);
        }

        let choices = super::RandomChoice::random_choice_f64(&samples, &weights, 4 as usize);

        let mut weight_counter = HashMap::with_capacity(capacity);

        for choice in choices {
            let counter = weight_counter.entry(choice).or_insert(0);
            *counter += 1;
        }

        for key in weight_counter.keys() {
            println!("{}", key);
        }
    }

    #[test]
    fn test_random_choice_in_place_64() {
        let mut samples = vec!["hi", "this", "is", "a", "test!"];
        let weights: Vec<f64> = vec![1.0, 1.0, 1.0, 1.0, 1.0];

        super::RandomChoice::random_choice_in_place_f64(&mut samples, &weights);

        for sample in samples {
            print!("{}, ", sample);
        }
    }

}<|endoftext|>