<commit_before>use std::fmt;
use std::mem;
use std::marker::PhantomData;
use std::net::SocketAddr;

use futures::{Async, Future, Poll, Stream};
use futures::stream::Concat2;
use hyper::{HeaderMap, StatusCode, Version};
use hyper::client::connect::HttpInfo;
use hyper::header::{CONTENT_LENGTH};
use serde::de::DeserializeOwned;
use serde_json;
use url::Url;
use http;

use super::Decoder;
use super::body::Body;

/// A Response to a submitted `Request`.
pub struct Response {
    status: StatusCode,
    headers: HeaderMap,
    // Boxed to save space (11 words to 1 word), and it's not accessed
    // frequently internally.
    url: Box<Url>,
    body: Decoder,
    version: Version,
    extensions: http::Extensions,
}

impl Response {
    pub(super) fn new(res: ::hyper::Response<::hyper::Body>, url: Url, gzip: bool) -> Response {
        let (parts, body) = res.into_parts();
        let status = parts.status;
        let version = parts.version;
        let extensions = parts.extensions;

        let mut headers = parts.headers;
        let decoder = Decoder::detect(&mut headers, Body::wrap(body), gzip);

        debug!("Response: '{}' for {}", status, url);
        Response {
            status,
            headers,
            url: Box::new(url),
            body: decoder,
            version,
            extensions,
        }
    }

    /// Get the `StatusCode` of this `Response`.
    #[inline]
    pub fn status(&self) -> StatusCode {
        self.status
    }

    /// Get the `Headers` of this `Response`.
    #[inline]
    pub fn headers(&self) -> &HeaderMap {
        &self.headers
    }

    /// Get a mutable reference to the `Headers` of this `Response`.
    #[inline]
    pub fn headers_mut(&mut self) -> &mut HeaderMap {
        &mut self.headers
    }

    /// Get the final `Url` of this `Response`.
    #[inline]
    pub fn url(&self) -> &Url {
        &self.url
    }

    /// Get the remote address used to get this `Response`.
    pub fn remote_addr(&self) -> Option<SocketAddr> {
        self
            .extensions
            .get::<HttpInfo>()
            .map(|info| info.remote_addr())
    }

    /// Get the content-length of this response, if known.
    ///
    /// Reasons it may not be known:
    ///
    /// - The server didn't send a `content-length` header.
    /// - The response is gzipped and automatically decoded (thus changing
    ///   the actual decoded length).
    pub fn content_length(&self) -> Option<u64> {
        self
            .headers()
            .get(CONTENT_LENGTH)
            .and_then(|ct_len| ct_len.to_str().ok())
            .and_then(|ct_len| ct_len.parse().ok())
    }

    /// Consumes the response, returning the body
    pub fn into_body(self) -> Decoder {
        self.body
    }

    /// Get a reference to the response body.
    #[inline]
    pub fn body(&self) -> &Decoder {
        &self.body
    }

    /// Get a mutable reference to the response body.
    ///
    /// The chunks from the body may be decoded, depending on the `gzip`
    /// option on the `ClientBuilder`.
    #[inline]
    pub fn body_mut(&mut self) -> &mut Decoder {
        &mut self.body
    }

    /// Get the HTTP `Version` of this `Response`.
    #[inline]
    pub fn version(&self) -> Version {
        self.version
    }

    /// Try to deserialize the response body as JSON using `serde`.
    #[inline]
    pub fn json<T: DeserializeOwned>(&mut self) -> Json<T> {
        let body = mem::replace(&mut self.body, Decoder::empty());

        Json {
            concat: body.concat2(),
            _marker: PhantomData,
        }
    }

    /// Turn a response into an error if the server returned an error.
    ///
    /// # Example
    ///
    /// ```
    /// # use reqwest::async::Response;
    /// fn on_response(res: Response) {
    ///     match res.error_for_status() {
    ///         Ok(_res) => (),
    ///         Err(err) => {
    ///             // asserting a 400 as an example
    ///             // it could be any status between 400...599
    ///             assert_eq!(
    ///                 err.status(),
    ///                 Some(reqwest::StatusCode::BAD_REQUEST)
    ///             );
    ///         }
    ///     }
    /// }
    /// # fn main() {}
    /// ```
    #[inline]
    pub fn error_for_status(self) -> ::Result<Self> {
        if self.status.is_client_error() {
            Err(::error::client_error(*self.url, self.status))
        } else if self.status.is_server_error() {
            Err(::error::server_error(*self.url, self.status))
        } else {
            Ok(self)
        }
    }
}

impl fmt::Debug for Response {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Response")
            .field("url", self.url())
            .field("status", &self.status())
            .field("headers", self.headers())
            .finish()
    }
}

impl<T: Into<Body>> From<http::Response<T>> for Response {
    fn from(r: http::Response<T>) -> Response {
        let (mut parts, body) = r.into_parts();
        let body = body.into();
        let body = Decoder::detect(&mut parts.headers, body, false);
        let url = parts.extensions
            .remove::<ResponseUrl>()
            .unwrap_or_else(|| ResponseUrl(Url::parse("http://no.url.provided.local").unwrap()));
        let url = url.0;
        Response {
            status: parts.status,
            headers: parts.headers,
            url: Box::new(url),
            body: body,
            version: parts.version,
            extensions: parts.extensions,
        }
    }
}

/// A JSON object.
pub struct Json<T> {
    concat: Concat2<Decoder>,
    _marker: PhantomData<T>,
}

impl<T: DeserializeOwned> Future for Json<T> {
    type Item = T;
    type Error = ::Error;
    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        let bytes = try_ready!(self.concat.poll());
        let t = try_!(serde_json::from_slice(&bytes));
        Ok(Async::Ready(t))
    }
}

impl<T> fmt::Debug for Json<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Json")
            .finish()
    }
}

#[derive(Debug, Clone, PartialEq)]
struct ResponseUrl(Url);

/// Extension trait for http::response::Builder objects
///
/// Allows the user to add a `Url` to the http::Response
pub trait ResponseBuilderExt {
    /// A builder method for the `http::response::Builder` type that allows the user to add a `Url`
    /// to the `http::Response`
    ///
    /// # Example
    ///
    /// ```
    /// # extern crate url;
    /// # extern crate http;
    /// # extern crate reqwest;
    /// # use std::error::Error;
    /// use url::Url;
    /// use http::response::Builder;
    /// use reqwest::async::ResponseBuilderExt;
    /// # fn main() -> Result<(), Box<Error>> {
    /// let response = Builder::new()
    ///     .status(200)
    ///     .url(Url::parse("http://example.com")?)
    ///     .body(())?;
    ///
    /// #   Ok(())
    /// # }
    fn url(&mut self, url: Url) -> &mut Self;
}

impl ResponseBuilderExt for http::response::Builder {
    fn url(&mut self, url: Url) -> &mut Self {
        self.extension(ResponseUrl(url))
    }
}

#[cfg(test)]
mod tests {
    use url::Url;
    use http::response::Builder;
    use super::{Response, ResponseUrl, ResponseBuilderExt};

    #[test]
    fn test_response_builder_ext() {
        let url = Url::parse("http://example.com").unwrap();
        let response = Builder::new()
            .status(200)
            .url(url.clone())
            .body(())
            .unwrap();

        assert_eq!(response.extensions().get::<ResponseUrl>(), Some(&ResponseUrl(url)));
    }

    #[test]
    fn test_from_http_response() {
        let url = Url::parse("http://example.com").unwrap();
        let response = Builder::new()
            .status(200)
            .url(url.clone())
            .body("foo")
            .unwrap();
        let response = Response::from(response);

        assert_eq!(response.status, 200);
        assert_eq!(response.url, Box::new(url));
    }
}
<commit_msg>async/reponse: return a impl Future on json()<commit_after>use std::fmt;
use std::mem;
use std::marker::PhantomData;
use std::net::SocketAddr;

use futures::{Async, Future, Poll, Stream};
use futures::stream::Concat2;
use hyper::{HeaderMap, StatusCode, Version};
use hyper::client::connect::HttpInfo;
use hyper::header::{CONTENT_LENGTH};
use serde::de::DeserializeOwned;
use serde_json;
use url::Url;
use http;

use super::Decoder;
use super::body::Body;

/// A Response to a submitted `Request`.
pub struct Response {
    status: StatusCode,
    headers: HeaderMap,
    // Boxed to save space (11 words to 1 word), and it's not accessed
    // frequently internally.
    url: Box<Url>,
    body: Decoder,
    version: Version,
    extensions: http::Extensions,
}

impl Response {
    pub(super) fn new(res: ::hyper::Response<::hyper::Body>, url: Url, gzip: bool) -> Response {
        let (parts, body) = res.into_parts();
        let status = parts.status;
        let version = parts.version;
        let extensions = parts.extensions;

        let mut headers = parts.headers;
        let decoder = Decoder::detect(&mut headers, Body::wrap(body), gzip);

        debug!("Response: '{}' for {}", status, url);
        Response {
            status,
            headers,
            url: Box::new(url),
            body: decoder,
            version,
            extensions,
        }
    }

    /// Get the `StatusCode` of this `Response`.
    #[inline]
    pub fn status(&self) -> StatusCode {
        self.status
    }

    /// Get the `Headers` of this `Response`.
    #[inline]
    pub fn headers(&self) -> &HeaderMap {
        &self.headers
    }

    /// Get a mutable reference to the `Headers` of this `Response`.
    #[inline]
    pub fn headers_mut(&mut self) -> &mut HeaderMap {
        &mut self.headers
    }

    /// Get the final `Url` of this `Response`.
    #[inline]
    pub fn url(&self) -> &Url {
        &self.url
    }

    /// Get the remote address used to get this `Response`.
    pub fn remote_addr(&self) -> Option<SocketAddr> {
        self
            .extensions
            .get::<HttpInfo>()
            .map(|info| info.remote_addr())
    }

    /// Get the content-length of this response, if known.
    ///
    /// Reasons it may not be known:
    ///
    /// - The server didn't send a `content-length` header.
    /// - The response is gzipped and automatically decoded (thus changing
    ///   the actual decoded length).
    pub fn content_length(&self) -> Option<u64> {
        self
            .headers()
            .get(CONTENT_LENGTH)
            .and_then(|ct_len| ct_len.to_str().ok())
            .and_then(|ct_len| ct_len.parse().ok())
    }

    /// Consumes the response, returning the body
    pub fn into_body(self) -> Decoder {
        self.body
    }

    /// Get a reference to the response body.
    #[inline]
    pub fn body(&self) -> &Decoder {
        &self.body
    }

    /// Get a mutable reference to the response body.
    ///
    /// The chunks from the body may be decoded, depending on the `gzip`
    /// option on the `ClientBuilder`.
    #[inline]
    pub fn body_mut(&mut self) -> &mut Decoder {
        &mut self.body
    }

    /// Get the HTTP `Version` of this `Response`.
    #[inline]
    pub fn version(&self) -> Version {
        self.version
    }

    /// Try to deserialize the response body as JSON using `serde`.
    #[inline]
    pub fn json<T: DeserializeOwned>(&mut self) -> impl Future<Item = T, Error = ::Error> {
        let body = mem::replace(&mut self.body, Decoder::empty());

        Json {
            concat: body.concat2(),
            _marker: PhantomData,
        }
    }

    /// Turn a response into an error if the server returned an error.
    ///
    /// # Example
    ///
    /// ```
    /// # use reqwest::async::Response;
    /// fn on_response(res: Response) {
    ///     match res.error_for_status() {
    ///         Ok(_res) => (),
    ///         Err(err) => {
    ///             // asserting a 400 as an example
    ///             // it could be any status between 400...599
    ///             assert_eq!(
    ///                 err.status(),
    ///                 Some(reqwest::StatusCode::BAD_REQUEST)
    ///             );
    ///         }
    ///     }
    /// }
    /// # fn main() {}
    /// ```
    #[inline]
    pub fn error_for_status(self) -> ::Result<Self> {
        if self.status.is_client_error() {
            Err(::error::client_error(*self.url, self.status))
        } else if self.status.is_server_error() {
            Err(::error::server_error(*self.url, self.status))
        } else {
            Ok(self)
        }
    }
}

impl fmt::Debug for Response {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Response")
            .field("url", self.url())
            .field("status", &self.status())
            .field("headers", self.headers())
            .finish()
    }
}

impl<T: Into<Body>> From<http::Response<T>> for Response {
    fn from(r: http::Response<T>) -> Response {
        let (mut parts, body) = r.into_parts();
        let body = body.into();
        let body = Decoder::detect(&mut parts.headers, body, false);
        let url = parts.extensions
            .remove::<ResponseUrl>()
            .unwrap_or_else(|| ResponseUrl(Url::parse("http://no.url.provided.local").unwrap()));
        let url = url.0;
        Response {
            status: parts.status,
            headers: parts.headers,
            url: Box::new(url),
            body: body,
            version: parts.version,
            extensions: parts.extensions,
        }
    }
}

/// A JSON object.
pub struct Json<T> {
    concat: Concat2<Decoder>,
    _marker: PhantomData<T>,
}

impl<T: DeserializeOwned> Future for Json<T> {
    type Item = T;
    type Error = ::Error;
    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        let bytes = try_ready!(self.concat.poll());
        let t = try_!(serde_json::from_slice(&bytes));
        Ok(Async::Ready(t))
    }
}

impl<T> fmt::Debug for Json<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Json")
            .finish()
    }
}

#[derive(Debug, Clone, PartialEq)]
struct ResponseUrl(Url);

/// Extension trait for http::response::Builder objects
///
/// Allows the user to add a `Url` to the http::Response
pub trait ResponseBuilderExt {
    /// A builder method for the `http::response::Builder` type that allows the user to add a `Url`
    /// to the `http::Response`
    ///
    /// # Example
    ///
    /// ```
    /// # extern crate url;
    /// # extern crate http;
    /// # extern crate reqwest;
    /// # use std::error::Error;
    /// use url::Url;
    /// use http::response::Builder;
    /// use reqwest::async::ResponseBuilderExt;
    /// # fn main() -> Result<(), Box<Error>> {
    /// let response = Builder::new()
    ///     .status(200)
    ///     .url(Url::parse("http://example.com")?)
    ///     .body(())?;
    ///
    /// #   Ok(())
    /// # }
    fn url(&mut self, url: Url) -> &mut Self;
}

impl ResponseBuilderExt for http::response::Builder {
    fn url(&mut self, url: Url) -> &mut Self {
        self.extension(ResponseUrl(url))
    }
}

#[cfg(test)]
mod tests {
    use url::Url;
    use http::response::Builder;
    use super::{Response, ResponseUrl, ResponseBuilderExt};

    #[test]
    fn test_response_builder_ext() {
        let url = Url::parse("http://example.com").unwrap();
        let response = Builder::new()
            .status(200)
            .url(url.clone())
            .body(())
            .unwrap();

        assert_eq!(response.extensions().get::<ResponseUrl>(), Some(&ResponseUrl(url)));
    }

    #[test]
    fn test_from_http_response() {
        let url = Url::parse("http://example.com").unwrap();
        let response = Builder::new()
            .status(200)
            .url(url.clone())
            .body("foo")
            .unwrap();
        let response = Response::from(response);

        assert_eq!(response.status, 200);
        assert_eq!(response.url, Box::new(url));
    }
}
<|endoftext|>


<commit_before>use std::cell::UnsafeCell;
use std::ptr;

use {Id, ShareId, Message};
use runtime::Object;

#[link(name = "objc", kind = "dylib")]
extern {
    fn objc_storeWeak(location: *mut *mut Object, obj: *mut Object) -> *mut Object;
    fn objc_loadWeakRetained(location: *mut *mut Object) -> *mut Object;
    fn objc_destroyWeak(addr: *mut *mut Object);
}

/// A pointer type for a weak reference to an Objective-C reference counted
/// object.
pub struct WeakId<T> {
    // Our pointer must have the same address even if we are moved, so Box it.
    // Although loading the WeakId may modify the pointer, it is thread safe,
    // so we must use an UnsafeCell to get a *mut without self being mutable.
    ptr: Box<UnsafeCell<*const T>>,
}

impl<T> WeakId<T> where T: Message {
    /// Construct a new `WeakId` referencing the given `ShareId`.
    pub fn new(obj: &ShareId<T>) -> WeakId<T> {
        let ptr = Box::new(UnsafeCell::new(ptr::null()));
        unsafe {
            let loc = ptr.get() as *mut *mut Object;
            objc_storeWeak(loc, &**obj as *const T as *mut Object);
        }
        WeakId { ptr: ptr }
    }

    /// Load a `ShareId` from the `WeakId` if the object still exists.
    /// Returns `None` if the object has been deallocated.
    pub fn load(&self) -> Option<ShareId<T>> {
        unsafe {
            let loc = self.ptr.get() as *mut *mut Object;
            let obj = objc_loadWeakRetained(loc) as *mut T;
            if obj.is_null() {
                None
            } else {
                Some(Id::from_retained_ptr(obj))
            }
        }
    }
}

#[unsafe_destructor]
impl<T> Drop for WeakId<T> {
    fn drop(&mut self) {
        unsafe {
            let loc = self.ptr.get() as *mut *mut Object;
            objc_destroyWeak(loc);
        }
    }
}

unsafe impl<T> Sync for WeakId<T> where T: Sync { }

unsafe impl<T> Send for WeakId<T> where T: Sync { }

#[cfg(test)]
mod tests {
    use runtime::Object;
    use test_utils;
    use super::WeakId;

    #[test]
    fn test_weak() {
        let obj = test_utils::sample_object().share();

        let weak = WeakId::new(&obj);
        let strong = weak.load().unwrap();
        assert!(&*strong as *const Object == &*obj as *const Object);
        drop(strong);

        drop(obj);
        assert!(weak.load().is_none());
    }
}
<commit_msg>Added a WeakPtr struct to remove unsafe_destructor.<commit_after>use std::cell::UnsafeCell;
use std::marker::PhantomData;
use std::ptr;

use {Id, ShareId, Message};
use runtime::Object;

#[link(name = "objc", kind = "dylib")]
extern {
    fn objc_storeWeak(location: *mut *mut Object, obj: *mut Object) -> *mut Object;
    fn objc_loadWeakRetained(location: *mut *mut Object) -> *mut Object;
    fn objc_destroyWeak(addr: *mut *mut Object);
}

// Our pointer must have the same address even if we are moved, so Box it.
// Although loading the WeakPtr may modify the pointer, it is thread safe,
// so we must use an UnsafeCell to get a *mut without self being mutable.
struct WeakPtr(Box<UnsafeCell<*mut Object>>);

impl WeakPtr {
    fn new(obj: *mut Object) -> WeakPtr {
        let ptr = Box::new(UnsafeCell::new(ptr::null_mut()));
        unsafe {
            objc_storeWeak(ptr.get(), obj);
        }
        WeakPtr(ptr)
    }

    fn load(&self) -> *mut Object {
        unsafe {
            objc_loadWeakRetained(self.0.get())
        }
    }
}

impl Drop for WeakPtr {
    fn drop(&mut self) {
        unsafe {
            objc_destroyWeak(self.0.get());
        }
    }
}

/// A pointer type for a weak reference to an Objective-C reference counted
/// object.
pub struct WeakId<T> {
    ptr: WeakPtr,
    item: PhantomData<T>,
}

impl<T> WeakId<T> where T: Message {
    /// Construct a new `WeakId` referencing the given `ShareId`.
    pub fn new(obj: &ShareId<T>) -> WeakId<T> {
        WeakId {
            ptr: WeakPtr::new(&**obj as *const T as *mut Object),
            item: PhantomData,
        }
    }

    /// Load a `ShareId` from the `WeakId` if the object still exists.
    /// Returns `None` if the object has been deallocated.
    pub fn load(&self) -> Option<ShareId<T>> {
        let obj = self.ptr.load();
        if obj.is_null() {
            None
        } else {
            unsafe { Some(Id::from_retained_ptr(obj as *mut T)) }
        }
    }
}

unsafe impl<T> Sync for WeakId<T> where T: Sync { }

unsafe impl<T> Send for WeakId<T> where T: Sync { }

#[cfg(test)]
mod tests {
    use runtime::Object;
    use test_utils;
    use super::WeakId;

    #[test]
    fn test_weak() {
        let obj = test_utils::sample_object().share();

        let weak = WeakId::new(&obj);
        let strong = weak.load().unwrap();
        assert!(&*strong as *const Object == &*obj as *const Object);
        drop(strong);

        drop(obj);
        assert!(weak.load().is_none());
    }
}
<|endoftext|>


<commit_before><commit_msg>screen: remove unsafe<commit_after><|endoftext|>


<commit_before>use serde_json::builder::ObjectBuilder;
use serde_json::Value;
use std::net::Shutdown;
use std::sync::mpsc::{
    TryRecvError,
    Receiver as MpscReceiver,
    Sender as MpscSender
};
use std::time::Duration as StdDuration;
use std::{env, thread};
use super::super::ClientError;
use super::{GatewayError, GatewayStatus};
use time::{self, Duration};
use websocket::client::request::Url as RequestUrl;
use websocket::client::{Receiver, Sender};
use websocket::stream::WebSocketStream;
use ::constants::{self, OpCode};
use ::error::{Error, Result};
use ::internal::ws_impl::{ReceiverExt, SenderExt};
use ::model::event::{Event, GatewayEvent, ReadyEvent};

#[inline]
pub fn parse_ready(event: GatewayEvent,
                   tx: &MpscSender<GatewayStatus>,
                   receiver: &mut Receiver<WebSocketStream>,
                   identification: Value)
                   -> Result<(ReadyEvent, u64)> {
    match event {
        GatewayEvent::Dispatch(seq, Event::Ready(event)) => {
            Ok((event, seq))
        },
        GatewayEvent::InvalidateSession => {
            debug!("Session invalidation");

            let _ = tx.send(GatewayStatus::SendMessage(identification));

            match receiver.recv_json(GatewayEvent::decode)? {
                GatewayEvent::Dispatch(seq, Event::Ready(event)) => {
                    Ok((event, seq))
                },
                other => {
                    debug!("Unexpected event: {:?}", other);

                    Err(Error::Gateway(GatewayError::InvalidHandshake))
                },
            }
        },
        other => {
            debug!("Unexpected event: {:?}", other);

            Err(Error::Gateway(GatewayError::InvalidHandshake))
        },
    }
}

pub fn identify(token: &str, shard_info: Option<[u8; 2]>) -> Value {
    ObjectBuilder::new()
        .insert("op", OpCode::Identify.num())
        .insert_object("d", |mut object| {
            object = identify_compression(object)
                .insert("large_threshold", 250) // max value
                .insert_object("properties", |object| object
                    .insert("$browser", "Ergonomic and high-level Rust library")
                    .insert("$device", "serenity")
                    .insert("$os", env::consts::OS))
                .insert("token", token)
                .insert("v", constants::GATEWAY_VERSION);

            if let Some(shard_info) = shard_info {
                object = object.insert_array("shard", |a| a
                    .push(shard_info[0])
                    .push(shard_info[1]));
            }

            object
        })
        .build()
}

#[cfg(not(feature="debug"))]
pub fn identify_compression(object: ObjectBuilder) -> ObjectBuilder {
    object.insert("compression", true)
}

#[cfg(feature="debug")]
pub fn identify_compression(object: ObjectBuilder) -> ObjectBuilder {
    object.insert("compression", false)
}

pub fn build_gateway_url(base: &str) -> Result<RequestUrl> {
    RequestUrl::parse(&format!("{}?v={}", base, constants::GATEWAY_VERSION))
        .map_err(|_| Error::Client(ClientError::Gateway))
}

pub fn keepalive(interval: u64,
                 mut sender: Sender<WebSocketStream>,
                 channel: MpscReceiver<GatewayStatus>) {
    let mut base_interval = Duration::milliseconds(interval as i64);
    let mut next_tick = time::get_time() + base_interval;

    let mut last_sequence = 0;

    'outer: loop {
        thread::sleep(StdDuration::from_millis(100));

        loop {
            match channel.try_recv() {
                Ok(GatewayStatus::ChangeInterval(interval)) => {
                    base_interval = Duration::milliseconds(interval as i64);
                },
                Ok(GatewayStatus::ChangeSender(new_sender)) => {
                    sender = new_sender;
                },
                Ok(GatewayStatus::SendMessage(val)) => {
                    if let Err(why) = sender.send_json(&val) {
                        warn!("Error sending message: {:?}", why);
                    }
                },
                Ok(GatewayStatus::Sequence(seq)) => {
                    last_sequence = seq;
                },
                Err(TryRecvError::Empty) => break,
                Err(TryRecvError::Disconnected) => break 'outer,
            }
        }

        if time::get_time() >= next_tick {
            next_tick = next_tick + base_interval;

            let map = ObjectBuilder::new()
                .insert("d", last_sequence)
                .insert("op", OpCode::Heartbeat.num())
                .build();

            if let Err(why) = sender.send_json(&map) {
                warn!("Error sending keepalive: {:?}", why);
            }
        }
    }

    let _ = sender.get_mut().shutdown(Shutdown::Both);
}
<commit_msg>Simplify gateway identify compression<commit_after>use serde_json::builder::ObjectBuilder;
use serde_json::Value;
use std::net::Shutdown;
use std::sync::mpsc::{
    TryRecvError,
    Receiver as MpscReceiver,
    Sender as MpscSender
};
use std::time::Duration as StdDuration;
use std::{env, thread};
use super::super::ClientError;
use super::{GatewayError, GatewayStatus};
use time::{self, Duration};
use websocket::client::request::Url as RequestUrl;
use websocket::client::{Receiver, Sender};
use websocket::stream::WebSocketStream;
use ::constants::{self, OpCode};
use ::error::{Error, Result};
use ::internal::ws_impl::{ReceiverExt, SenderExt};
use ::model::event::{Event, GatewayEvent, ReadyEvent};

#[inline]
pub fn parse_ready(event: GatewayEvent,
                   tx: &MpscSender<GatewayStatus>,
                   receiver: &mut Receiver<WebSocketStream>,
                   identification: Value)
                   -> Result<(ReadyEvent, u64)> {
    match event {
        GatewayEvent::Dispatch(seq, Event::Ready(event)) => {
            Ok((event, seq))
        },
        GatewayEvent::InvalidateSession => {
            debug!("Session invalidation");

            let _ = tx.send(GatewayStatus::SendMessage(identification));

            match receiver.recv_json(GatewayEvent::decode)? {
                GatewayEvent::Dispatch(seq, Event::Ready(event)) => {
                    Ok((event, seq))
                },
                other => {
                    debug!("Unexpected event: {:?}", other);

                    Err(Error::Gateway(GatewayError::InvalidHandshake))
                },
            }
        },
        other => {
            debug!("Unexpected event: {:?}", other);

            Err(Error::Gateway(GatewayError::InvalidHandshake))
        },
    }
}

pub fn identify(token: &str, shard_info: Option<[u8; 2]>) -> Value {
    ObjectBuilder::new()
        .insert("op", OpCode::Identify.num())
        .insert_object("d", |mut object| {
            object = identify_compression(object)
                .insert("large_threshold", 250) // max value
                .insert_object("properties", |object| object
                    .insert("$browser", "Ergonomic and high-level Rust library")
                    .insert("$device", "serenity")
                    .insert("$os", env::consts::OS))
                .insert("token", token)
                .insert("v", constants::GATEWAY_VERSION);

            if let Some(shard_info) = shard_info {
                object = object.insert_array("shard", |a| a
                    .push(shard_info[0])
                    .push(shard_info[1]));
            }

            object
        })
        .build()
}

#[inline(always)]
pub fn identify_compression(object: ObjectBuilder) -> ObjectBuilder {
    object.insert("compression", !cfg!(feature="debug"))
}

pub fn build_gateway_url(base: &str) -> Result<RequestUrl> {
    RequestUrl::parse(&format!("{}?v={}", base, constants::GATEWAY_VERSION))
        .map_err(|_| Error::Client(ClientError::Gateway))
}

pub fn keepalive(interval: u64,
                 mut sender: Sender<WebSocketStream>,
                 channel: MpscReceiver<GatewayStatus>) {
    let mut base_interval = Duration::milliseconds(interval as i64);
    let mut next_tick = time::get_time() + base_interval;

    let mut last_sequence = 0;

    'outer: loop {
        thread::sleep(StdDuration::from_millis(100));

        loop {
            match channel.try_recv() {
                Ok(GatewayStatus::ChangeInterval(interval)) => {
                    base_interval = Duration::milliseconds(interval as i64);
                },
                Ok(GatewayStatus::ChangeSender(new_sender)) => {
                    sender = new_sender;
                },
                Ok(GatewayStatus::SendMessage(val)) => {
                    if let Err(why) = sender.send_json(&val) {
                        warn!("Error sending message: {:?}", why);
                    }
                },
                Ok(GatewayStatus::Sequence(seq)) => {
                    last_sequence = seq;
                },
                Err(TryRecvError::Empty) => break,
                Err(TryRecvError::Disconnected) => break 'outer,
            }
        }

        if time::get_time() >= next_tick {
            next_tick = next_tick + base_interval;

            let map = ObjectBuilder::new()
                .insert("d", last_sequence)
                .insert("op", OpCode::Heartbeat.num())
                .build();

            if let Err(why) = sender.send_json(&map) {
                warn!("Error sending keepalive: {:?}", why);
            }
        }
    }

    let _ = sender.get_mut().shutdown(Shutdown::Both);
}
<|endoftext|>


<commit_before>// This module defines a container that expresses zakocmp policy
// violations - i.e. human-readable descriptions of notable
// discrepancies between zakocmp snapshots.

use crate::structs::ZakocmpError;

// Represents possible policy violations as caller-passable ints.
pub const ADDED: i32 = 0;
pub const DELETED: i32 = 1 << 0;
pub const MODIFIED: i32 = 1 << 1;

// Represents human-readable counterparts to the above. These are meant
// for printing etc. and so are not relevant to callers.
const REPR_ADDED: &'static str = "+";
const REPR_DELETED: &'static str = "-";
const REPR_MODIFIED: &'static str = "!";

// A single violation consists of the offending path (arbitrary bytes)
// and the kind of violation (i32 - as above).
pub struct Violations {
    data: std::vec::Vec<(String, i32)>,
}

impl Violations {
    pub fn new() -> Violations {
        Violations { data: vec![] }
    }

    // Accepts a path (arbitrary string of bytes) and a kind of policy
    // violation.
    pub fn add(&mut self, path: &str, kind: i32) -> std::result::Result<(), ZakocmpError> {
        match kind {
            ADDED | DELETED | MODIFIED => (),
            _ => return Err(ZakocmpError::Unknown(format!("bad kind: {}", kind))),
        };
        self.data.push((path.to_owned(), kind));
        Ok(())
    }
}

fn display_violation_type(kind: i32) -> &'static str {
    match kind {
        ADDED => REPR_ADDED,
        DELETED => REPR_DELETED,
        MODIFIED => REPR_MODIFIED,
        // This case is serious: the burden is on us to have weeded out
        // invalid violation kinds in previous calls to add().
        _ => panic!(format!("BUG: bad kind: {}", kind)),
    }
}

impl std::fmt::Display for Violations {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut sorted_violations = self.data.to_owned();
        sorted_violations.sort_unstable();
        for (path, kind) in sorted_violations.into_iter() {
            write!(f, "{} {}\n", display_violation_type(kind), path)?;
        }
        Ok(())
    }
}

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

    #[test]
    fn violations_add_fail() {
        let mut violations = Violations::new();
        let error = violations.add("some/path", 999).unwrap_err();
        match error {
            ZakocmpError::Unknown(message) => assert!(message.starts_with("bad")),
            _ => panic!("not a ZakocmpError::Unknown"),
        };
    }

    #[test]
    fn violations_add_okay() {
        let mut violations = Violations::new();
        assert!(violations
            .add(
                "arbitrary bytes are allowed so long as the type is valid",
                ADDED
            )
            .is_ok());
        // Empty strings are allowed, too.
        assert!(violations.add("", MODIFIED).is_ok());
        // Doubling up on strings is perfectly legitimate, too.
        assert!(violations.add("", MODIFIED).is_ok());
        assert!(violations.add("", ADDED).is_ok());
        assert!(violations.add("", DELETED).is_ok());
    }

    #[test]
    fn violations_display() {
        let mut violations = Violations::new();
        assert!(violations.add("jello there!", ADDED).is_ok());
        assert!(violations.add("iello there!", MODIFIED).is_ok());
        assert!(violations.add("hello there!", DELETED).is_ok());
        assert!(violations.add("a/path/of/some/sort", ADDED).is_ok());
        assert!(violations.add("b/path/of/some/sort", MODIFIED).is_ok());
        assert!(violations.add("z/path/of/some/sort", DELETED).is_ok());

        assert_eq!(
            format!("{}", violations),
            r#"+ a/path/of/some/sort
! b/path/of/some/sort
- hello there!
! iello there!
+ jello there!
- z/path/of/some/sort
"#
        );
    }
}
<commit_msg>violations: use panic! formatter directly<commit_after>// This module defines a container that expresses zakocmp policy
// violations - i.e. human-readable descriptions of notable
// discrepancies between zakocmp snapshots.

use crate::structs::ZakocmpError;

// Represents possible policy violations as caller-passable ints.
pub const ADDED: i32 = 0;
pub const DELETED: i32 = 1 << 0;
pub const MODIFIED: i32 = 1 << 1;

// Represents human-readable counterparts to the above. These are meant
// for printing etc. and so are not relevant to callers.
const REPR_ADDED: &'static str = "+";
const REPR_DELETED: &'static str = "-";
const REPR_MODIFIED: &'static str = "!";

// A single violation consists of the offending path (arbitrary bytes)
// and the kind of violation (i32 - as above).
pub struct Violations {
    data: std::vec::Vec<(String, i32)>,
}

impl Violations {
    pub fn new() -> Violations {
        Violations { data: vec![] }
    }

    // Accepts a path (arbitrary string of bytes) and a kind of policy
    // violation.
    pub fn add(&mut self, path: &str, kind: i32) -> std::result::Result<(), ZakocmpError> {
        match kind {
            ADDED | DELETED | MODIFIED => (),
            _ => return Err(ZakocmpError::Unknown(format!("bad kind: {}", kind))),
        };
        self.data.push((path.to_owned(), kind));
        Ok(())
    }
}

fn display_violation_type(kind: i32) -> &'static str {
    match kind {
        ADDED => REPR_ADDED,
        DELETED => REPR_DELETED,
        MODIFIED => REPR_MODIFIED,
        // This case is serious: the burden is on us to have weeded out
        // invalid violation kinds in previous calls to add().
        _ => panic!("BUG: bad kind: {}", kind),
    }
}

impl std::fmt::Display for Violations {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut sorted_violations = self.data.to_owned();
        sorted_violations.sort_unstable();
        for (path, kind) in sorted_violations.into_iter() {
            write!(f, "{} {}\n", display_violation_type(kind), path)?;
        }
        Ok(())
    }
}

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

    #[test]
    fn violations_add_fail() {
        let mut violations = Violations::new();
        let error = violations.add("some/path", 999).unwrap_err();
        match error {
            ZakocmpError::Unknown(message) => assert!(message.starts_with("bad")),
            _ => panic!("not a ZakocmpError::Unknown"),
        };
    }

    #[test]
    fn violations_add_okay() {
        let mut violations = Violations::new();
        assert!(violations
            .add(
                "arbitrary bytes are allowed so long as the type is valid",
                ADDED
            )
            .is_ok());
        // Empty strings are allowed, too.
        assert!(violations.add("", MODIFIED).is_ok());
        // Doubling up on strings is perfectly legitimate, too.
        assert!(violations.add("", MODIFIED).is_ok());
        assert!(violations.add("", ADDED).is_ok());
        assert!(violations.add("", DELETED).is_ok());
    }

    #[test]
    fn violations_display() {
        let mut violations = Violations::new();
        assert!(violations.add("jello there!", ADDED).is_ok());
        assert!(violations.add("iello there!", MODIFIED).is_ok());
        assert!(violations.add("hello there!", DELETED).is_ok());
        assert!(violations.add("a/path/of/some/sort", ADDED).is_ok());
        assert!(violations.add("b/path/of/some/sort", MODIFIED).is_ok());
        assert!(violations.add("z/path/of/some/sort", DELETED).is_ok());

        assert_eq!(
            format!("{}", violations),
            r#"+ a/path/of/some/sort
! b/path/of/some/sort
- hello there!
! iello there!
+ jello there!
- z/path/of/some/sort
"#
        );
    }
}
<|endoftext|>


<commit_before>#![allow(dead_code)]

use ql2::types;
use types::string::IntoString;
use ql2::proto::Term_TermType as TermType;
use ::Client;
use serde_json::value::ToJson;

impl<O> Client<types::StreamSelection, O>
    where O: ToJson + Clone
{
    pub fn has_fields<T>(self, arg: T) -> Client<types::Stream, ()>
        where T: IntoString
        {
            super::make_cmd(TermType::GET_FIELD, Some(vec![arg.into_string()]), None, Some(self.cmd), self.errors)
        }
}

impl<O> Client<types::Object, O>
    where O: ToJson + Clone
{
    pub fn has_fields<T>(self, arg: T) -> Client<types::Bool, ()>
        where T: IntoString
        {
            super::make_cmd(TermType::GET_FIELD, Some(vec![arg.into_string()]), None, Some(self.cmd), self.errors)
        }
}

impl<O> Client<types::Array, O>
    where O: ToJson + Clone
{
    pub fn has_fields<T>(self, arg: T) -> Client<types::Array, ()>
        where T: IntoString
        {
            super::make_cmd(TermType::GET_FIELD, Some(vec![arg.into_string()]), None, Some(self.cmd), self.errors)
        }
}
<commit_msg>Fix has_fields term type<commit_after>#![allow(dead_code)]

use ql2::types;
use types::string::IntoString;
use ql2::proto::Term_TermType as TermType;
use ::Client;
use serde_json::value::ToJson;

impl<O> Client<types::StreamSelection, O>
    where O: ToJson + Clone
{
    pub fn has_fields<T>(self, arg: T) -> Client<types::Stream, ()>
        where T: IntoString
        {
            super::make_cmd(TermType::HAS_FIELDS, Some(vec![arg.into_string()]), None, Some(self.cmd), self.errors)
        }
}

impl<O> Client<types::Object, O>
    where O: ToJson + Clone
{
    pub fn has_fields<T>(self, arg: T) -> Client<types::Bool, ()>
        where T: IntoString
        {
            super::make_cmd(TermType::HAS_FIELDS, Some(vec![arg.into_string()]), None, Some(self.cmd), self.errors)
        }
}

impl<O> Client<types::Array, O>
    where O: ToJson + Clone
{
    pub fn has_fields<T>(self, arg: T) -> Client<types::Array, ()>
        where T: IntoString
        {
            super::make_cmd(TermType::HAS_FIELDS, Some(vec![arg.into_string()]), None, Some(self.cmd), self.errors)
        }
}
<|endoftext|>


<commit_before>

/*
 * The compiler code necessary to support the #fmt extension. Eventually this
 * should all get sucked into either the standard library extfmt module or the
 * compiler syntax extension plugin interface.
 */
import core::{vec, str, option};
import option::{some};
import extfmt::ct::*;
import base::*;
import codemap::span;
export expand_syntax_ext;

fn expand_syntax_ext(cx: ext_ctxt, sp: span, arg: @ast::expr,
                     _body: option::t<str>) -> @ast::expr {
    let args: [@ast::expr] =
        alt arg.node {
          ast::expr_vec(elts, _) { elts }
          _ {
            cx.span_fatal(sp, "#fmt requires arguments of the form `[...]`.")
          }
        };
    if vec::len::<@ast::expr>(args) == 0u {
        cx.span_fatal(sp, "#fmt requires a format string");
    }
    let fmt =
        expr_to_str(cx, args[0],
                    "first argument to #fmt must be a " + "string literal.");
    let fmtspan = args[0].span;
    #debug("Format string:");
    log(debug, fmt);
    fn parse_fmt_err_(cx: ext_ctxt, sp: span, msg: str) -> ! {
        cx.span_fatal(sp, msg);
    }
    let parse_fmt_err = bind parse_fmt_err_(cx, fmtspan, _);
    let pieces = parse_fmt_string(fmt, parse_fmt_err);
    ret pieces_to_expr(cx, sp, pieces, args);
}

// FIXME: A lot of these functions for producing expressions can probably
// be factored out in common with other code that builds expressions.
// FIXME: Cleanup the naming of these functions
fn pieces_to_expr(cx: ext_ctxt, sp: span, pieces: [piece], args: [@ast::expr])
   -> @ast::expr {
    fn make_new_lit(cx: ext_ctxt, sp: span, lit: ast::lit_) -> @ast::expr {
        let sp_lit = @{node: lit, span: sp};
        ret @{id: cx.next_id(), node: ast::expr_lit(sp_lit), span: sp};
    }
    fn make_new_str(cx: ext_ctxt, sp: span, s: str) -> @ast::expr {
        let lit = ast::lit_str(s);
        ret make_new_lit(cx, sp, lit);
    }
    fn make_new_int(cx: ext_ctxt, sp: span, i: int) -> @ast::expr {
        let lit = ast::lit_int(i as i64, ast::ty_i);
        ret make_new_lit(cx, sp, lit);
    }
    fn make_new_uint(cx: ext_ctxt, sp: span, u: uint) -> @ast::expr {
        let lit = ast::lit_uint(u as u64, ast::ty_u);
        ret make_new_lit(cx, sp, lit);
    }
    fn make_add_expr(cx: ext_ctxt, sp: span, lhs: @ast::expr, rhs: @ast::expr)
       -> @ast::expr {
        let binexpr = ast::expr_binary(ast::add, lhs, rhs);
        ret @{id: cx.next_id(), node: binexpr, span: sp};
    }
    fn make_path_expr(cx: ext_ctxt, sp: span, idents: [ast::ident]) ->
       @ast::expr {
        let path = {global: false, idents: idents, types: []};
        let sp_path = @{node: path, span: sp};
        let pathexpr = ast::expr_path(sp_path);
        ret @{id: cx.next_id(), node: pathexpr, span: sp};
    }
    fn make_vec_expr(cx: ext_ctxt, sp: span, exprs: [@ast::expr]) ->
       @ast::expr {
        let vecexpr = ast::expr_vec(exprs, ast::imm);
        ret @{id: cx.next_id(), node: vecexpr, span: sp};
    }
    fn make_call(cx: ext_ctxt, sp: span, fn_path: [ast::ident],
                 args: [@ast::expr]) -> @ast::expr {
        let pathexpr = make_path_expr(cx, sp, fn_path);
        let callexpr = ast::expr_call(pathexpr, args, false);
        ret @{id: cx.next_id(), node: callexpr, span: sp};
    }
    fn make_rec_expr(cx: ext_ctxt, sp: span,
                     fields: [{ident: ast::ident, ex: @ast::expr}]) ->
       @ast::expr {
        let astfields: [ast::field] = [];
        for field: {ident: ast::ident, ex: @ast::expr} in fields {
            let ident = field.ident;
            let val = field.ex;
            let astfield =
                {node: {mut: ast::imm, ident: ident, expr: val}, span: sp};
            astfields += [astfield];
        }
        let recexpr = ast::expr_rec(astfields, option::none::<@ast::expr>);
        ret @{id: cx.next_id(), node: recexpr, span: sp};
    }
    fn make_path_vec(_cx: ext_ctxt, ident: ast::ident) -> [ast::ident] {
        ret ["extfmt", "rt", ident];
    }
    fn make_rt_path_expr(cx: ext_ctxt, sp: span, ident: str) -> @ast::expr {
        let path = make_path_vec(cx, ident);
        ret make_path_expr(cx, sp, path);
    }
    // Produces an AST expression that represents a RT::conv record,
    // which tells the RT::conv* functions how to perform the conversion

    fn make_rt_conv_expr(cx: ext_ctxt, sp: span, cnv: conv) -> @ast::expr {
        fn make_flags(cx: ext_ctxt, sp: span, flags: [flag]) -> @ast::expr {
            let flagexprs: [@ast::expr] = [];
            for f: flag in flags {
                let fstr;
                alt f {
                  flag_left_justify. { fstr = "flag_left_justify"; }
                  flag_left_zero_pad. { fstr = "flag_left_zero_pad"; }
                  flag_space_for_sign. { fstr = "flag_space_for_sign"; }
                  flag_sign_always. { fstr = "flag_sign_always"; }
                  flag_alternate. { fstr = "flag_alternate"; }
                }
                flagexprs += [make_rt_path_expr(cx, sp, fstr)];
            }
            // FIXME: 0-length vectors can't have their type inferred
            // through the rec that these flags are a member of, so
            // this is a hack placeholder flag

            if vec::len::<@ast::expr>(flagexprs) == 0u {
                flagexprs += [make_rt_path_expr(cx, sp, "flag_none")];
            }
            ret make_vec_expr(cx, sp, flagexprs);
        }
        fn make_count(cx: ext_ctxt, sp: span, cnt: count) -> @ast::expr {
            alt cnt {
              count_implied. {
                ret make_rt_path_expr(cx, sp, "count_implied");
              }
              count_is(c) {
                let count_lit = make_new_int(cx, sp, c);
                let count_is_path = make_path_vec(cx, "count_is");
                let count_is_args = [count_lit];
                ret make_call(cx, sp, count_is_path, count_is_args);
              }
              _ { cx.span_unimpl(sp, "unimplemented #fmt conversion"); }
            }
        }
        fn make_ty(cx: ext_ctxt, sp: span, t: ty) -> @ast::expr {
            let rt_type;
            alt t {
              ty_hex(c) {
                alt c {
                  case_upper. { rt_type = "ty_hex_upper"; }
                  case_lower. { rt_type = "ty_hex_lower"; }
                }
              }
              ty_bits. { rt_type = "ty_bits"; }
              ty_octal. { rt_type = "ty_octal"; }
              _ { rt_type = "ty_default"; }
            }
            ret make_rt_path_expr(cx, sp, rt_type);
        }
        fn make_conv_rec(cx: ext_ctxt, sp: span, flags_expr: @ast::expr,
                         width_expr: @ast::expr, precision_expr: @ast::expr,
                         ty_expr: @ast::expr) -> @ast::expr {
            ret make_rec_expr(cx, sp,
                              [{ident: "flags", ex: flags_expr},
                               {ident: "width", ex: width_expr},
                               {ident: "precision", ex: precision_expr},
                               {ident: "ty", ex: ty_expr}]);
        }
        let rt_conv_flags = make_flags(cx, sp, cnv.flags);
        let rt_conv_width = make_count(cx, sp, cnv.width);
        let rt_conv_precision = make_count(cx, sp, cnv.precision);
        let rt_conv_ty = make_ty(cx, sp, cnv.ty);
        ret make_conv_rec(cx, sp, rt_conv_flags, rt_conv_width,
                          rt_conv_precision, rt_conv_ty);
    }
    fn make_conv_call(cx: ext_ctxt, sp: span, conv_type: str, cnv: conv,
                      arg: @ast::expr) -> @ast::expr {
        let fname = "conv_" + conv_type;
        let path = make_path_vec(cx, fname);
        let cnv_expr = make_rt_conv_expr(cx, sp, cnv);
        let args = [cnv_expr, arg];
        ret make_call(cx, arg.span, path, args);
    }
    fn make_new_conv(cx: ext_ctxt, sp: span, cnv: conv, arg: @ast::expr) ->
       @ast::expr {
        // FIXME: Extract all this validation into extfmt::ct

        fn is_signed_type(cnv: conv) -> bool {
            alt cnv.ty {
              ty_int(s) {
                alt s { signed. { ret true; } unsigned. { ret false; } }
              }
              ty_float. { ret true; }
              _ { ret false; }
            }
        }
        let unsupported = "conversion not supported in #fmt string";
        alt cnv.param {
          option::none. { }
          _ { cx.span_unimpl(sp, unsupported); }
        }
        for f: flag in cnv.flags {
            alt f {
              flag_left_justify. { }
              flag_sign_always. {
                if !is_signed_type(cnv) {
                    cx.span_fatal(sp,
                                  "+ flag only valid in " +
                                      "signed #fmt conversion");
                }
              }
              flag_space_for_sign. {
                if !is_signed_type(cnv) {
                    cx.span_fatal(sp,
                                  "space flag only valid in " +
                                      "signed #fmt conversions");
                }
              }
              flag_left_zero_pad. { }
              _ { cx.span_unimpl(sp, unsupported); }
            }
        }
        alt cnv.width {
          count_implied. { }
          count_is(_) { }
          _ { cx.span_unimpl(sp, unsupported); }
        }
        alt cnv.precision {
          count_implied. { }
          count_is(_) { }
          _ { cx.span_unimpl(sp, unsupported); }
        }
        alt cnv.ty {
          ty_str. { ret make_conv_call(cx, arg.span, "str", cnv, arg); }
          ty_int(sign) {
            alt sign {
              signed. { ret make_conv_call(cx, arg.span, "int", cnv, arg); }
              unsigned. {
                ret make_conv_call(cx, arg.span, "uint", cnv, arg);
              }
            }
          }
          ty_bool. { ret make_conv_call(cx, arg.span, "bool", cnv, arg); }
          ty_char. { ret make_conv_call(cx, arg.span, "char", cnv, arg); }
          ty_hex(_) { ret make_conv_call(cx, arg.span, "uint", cnv, arg); }
          ty_bits. { ret make_conv_call(cx, arg.span, "uint", cnv, arg); }
          ty_octal. { ret make_conv_call(cx, arg.span, "uint", cnv, arg); }
          ty_float. { ret make_conv_call(cx, arg.span, "float", cnv, arg); }
          _ { cx.span_unimpl(sp, unsupported); }
        }
    }
    fn log_conv(c: conv) {
        alt c.param {
          some(p) { log(debug, "param: " + int::to_str(p, 10u)); }
          _ { #debug("param: none"); }
        }
        for f: flag in c.flags {
            alt f {
              flag_left_justify. { #debug("flag: left justify"); }
              flag_left_zero_pad. { #debug("flag: left zero pad"); }
              flag_space_for_sign. { #debug("flag: left space pad"); }
              flag_sign_always. { #debug("flag: sign always"); }
              flag_alternate. { #debug("flag: alternate"); }
            }
        }
        alt c.width {
          count_is(i) { log(debug,
                                 "width: count is " + int::to_str(i, 10u)); }
          count_is_param(i) {
            log(debug,
                     "width: count is param " + int::to_str(i, 10u));
          }
          count_is_next_param. { #debug("width: count is next param"); }
          count_implied. { #debug("width: count is implied"); }
        }
        alt c.precision {
          count_is(i) { log(debug,
                                 "prec: count is " + int::to_str(i, 10u)); }
          count_is_param(i) {
            log(debug,
                     "prec: count is param " + int::to_str(i, 10u));
          }
          count_is_next_param. { #debug("prec: count is next param"); }
          count_implied. { #debug("prec: count is implied"); }
        }
        alt c.ty {
          ty_bool. { #debug("type: bool"); }
          ty_str. { #debug("type: str"); }
          ty_char. { #debug("type: char"); }
          ty_int(s) {
            alt s {
              signed. { #debug("type: signed"); }
              unsigned. { #debug("type: unsigned"); }
            }
          }
          ty_bits. { #debug("type: bits"); }
          ty_hex(cs) {
            alt cs {
              case_upper. { #debug("type: uhex"); }
              case_lower. { #debug("type: lhex"); }
            }
          }
          ty_octal. { #debug("type: octal"); }
          ty_float. { #debug("type: float"); }
        }
    }
    let fmt_sp = args[0].span;
    let n = 0u;
    let tmp_expr = make_new_str(cx, sp, "");
    let nargs = vec::len::<@ast::expr>(args);
    for pc: piece in pieces {
        alt pc {
          piece_string(s) {
            let s_expr = make_new_str(cx, fmt_sp, s);
            tmp_expr = make_add_expr(cx, fmt_sp, tmp_expr, s_expr);
          }
          piece_conv(conv) {
            n += 1u;
            if n >= nargs {
                cx.span_fatal(sp,
                              "not enough arguments to #fmt " +
                                  "for the given format string");
            }
            #debug("Building conversion:");
            log_conv(conv);
            let arg_expr = args[n];
            let c_expr = make_new_conv(cx, fmt_sp, conv, arg_expr);
            tmp_expr = make_add_expr(cx, fmt_sp, tmp_expr, c_expr);
          }
        }
    }
    let expected_nargs = n + 1u; // n conversions + the fmt string

    if expected_nargs < nargs {
        cx.span_fatal
            (sp, #fmt["too many arguments to #fmt. found %u, expected %u",
                           nargs, expected_nargs]);
    }
    ret tmp_expr;
}
//
// Local Variables:
// mode: rust
// fill-column: 78;
// indent-tabs-mode: nil
// c-basic-offset: 4
// buffer-file-coding-system: utf-8-unix
// End:
//
<commit_msg>comp: minor cleanup.<commit_after>

/*
 * The compiler code necessary to support the #fmt extension. Eventually this
 * should all get sucked into either the standard library extfmt module or the
 * compiler syntax extension plugin interface.
 */
import core::{vec, str, option};
import option::{some};
import extfmt::ct::*;
import base::*;
import codemap::span;
export expand_syntax_ext;

fn expand_syntax_ext(cx: ext_ctxt, sp: span, arg: @ast::expr,
                     _body: option::t<str>) -> @ast::expr {
    let args: [@ast::expr] =
        alt arg.node {
          ast::expr_vec(elts, _) { elts }
          _ {
            cx.span_fatal(sp, "#fmt requires arguments of the form `[...]`.")
          }
        };
    if vec::len::<@ast::expr>(args) == 0u {
        cx.span_fatal(sp, "#fmt requires a format string");
    }
    let fmt =
        expr_to_str(cx, args[0],
                    "first argument to #fmt must be a string literal.");
    let fmtspan = args[0].span;
    #debug("Format string:");
    log(debug, fmt);
    fn parse_fmt_err_(cx: ext_ctxt, sp: span, msg: str) -> ! {
        cx.span_fatal(sp, msg);
    }
    let parse_fmt_err = bind parse_fmt_err_(cx, fmtspan, _);
    let pieces = parse_fmt_string(fmt, parse_fmt_err);
    ret pieces_to_expr(cx, sp, pieces, args);
}

// FIXME: A lot of these functions for producing expressions can probably
// be factored out in common with other code that builds expressions.
// FIXME: Cleanup the naming of these functions
fn pieces_to_expr(cx: ext_ctxt, sp: span, pieces: [piece], args: [@ast::expr])
   -> @ast::expr {
    fn make_new_lit(cx: ext_ctxt, sp: span, lit: ast::lit_) -> @ast::expr {
        let sp_lit = @{node: lit, span: sp};
        ret @{id: cx.next_id(), node: ast::expr_lit(sp_lit), span: sp};
    }
    fn make_new_str(cx: ext_ctxt, sp: span, s: str) -> @ast::expr {
        let lit = ast::lit_str(s);
        ret make_new_lit(cx, sp, lit);
    }
    fn make_new_int(cx: ext_ctxt, sp: span, i: int) -> @ast::expr {
        let lit = ast::lit_int(i as i64, ast::ty_i);
        ret make_new_lit(cx, sp, lit);
    }
    fn make_new_uint(cx: ext_ctxt, sp: span, u: uint) -> @ast::expr {
        let lit = ast::lit_uint(u as u64, ast::ty_u);
        ret make_new_lit(cx, sp, lit);
    }
    fn make_add_expr(cx: ext_ctxt, sp: span, lhs: @ast::expr, rhs: @ast::expr)
       -> @ast::expr {
        let binexpr = ast::expr_binary(ast::add, lhs, rhs);
        ret @{id: cx.next_id(), node: binexpr, span: sp};
    }
    fn make_path_expr(cx: ext_ctxt, sp: span, idents: [ast::ident]) ->
       @ast::expr {
        let path = {global: false, idents: idents, types: []};
        let sp_path = @{node: path, span: sp};
        let pathexpr = ast::expr_path(sp_path);
        ret @{id: cx.next_id(), node: pathexpr, span: sp};
    }
    fn make_vec_expr(cx: ext_ctxt, sp: span, exprs: [@ast::expr]) ->
       @ast::expr {
        let vecexpr = ast::expr_vec(exprs, ast::imm);
        ret @{id: cx.next_id(), node: vecexpr, span: sp};
    }
    fn make_call(cx: ext_ctxt, sp: span, fn_path: [ast::ident],
                 args: [@ast::expr]) -> @ast::expr {
        let pathexpr = make_path_expr(cx, sp, fn_path);
        let callexpr = ast::expr_call(pathexpr, args, false);
        ret @{id: cx.next_id(), node: callexpr, span: sp};
    }
    fn make_rec_expr(cx: ext_ctxt, sp: span,
                     fields: [{ident: ast::ident, ex: @ast::expr}]) ->
       @ast::expr {
        let astfields: [ast::field] = [];
        for field: {ident: ast::ident, ex: @ast::expr} in fields {
            let ident = field.ident;
            let val = field.ex;
            let astfield =
                {node: {mut: ast::imm, ident: ident, expr: val}, span: sp};
            astfields += [astfield];
        }
        let recexpr = ast::expr_rec(astfields, option::none::<@ast::expr>);
        ret @{id: cx.next_id(), node: recexpr, span: sp};
    }
    fn make_path_vec(_cx: ext_ctxt, ident: ast::ident) -> [ast::ident] {
        ret ["extfmt", "rt", ident];
    }
    fn make_rt_path_expr(cx: ext_ctxt, sp: span, ident: str) -> @ast::expr {
        let path = make_path_vec(cx, ident);
        ret make_path_expr(cx, sp, path);
    }
    // Produces an AST expression that represents a RT::conv record,
    // which tells the RT::conv* functions how to perform the conversion

    fn make_rt_conv_expr(cx: ext_ctxt, sp: span, cnv: conv) -> @ast::expr {
        fn make_flags(cx: ext_ctxt, sp: span, flags: [flag]) -> @ast::expr {
            let flagexprs: [@ast::expr] = [];
            for f: flag in flags {
                let fstr;
                alt f {
                  flag_left_justify. { fstr = "flag_left_justify"; }
                  flag_left_zero_pad. { fstr = "flag_left_zero_pad"; }
                  flag_space_for_sign. { fstr = "flag_space_for_sign"; }
                  flag_sign_always. { fstr = "flag_sign_always"; }
                  flag_alternate. { fstr = "flag_alternate"; }
                }
                flagexprs += [make_rt_path_expr(cx, sp, fstr)];
            }
            // FIXME: 0-length vectors can't have their type inferred
            // through the rec that these flags are a member of, so
            // this is a hack placeholder flag

            if vec::len::<@ast::expr>(flagexprs) == 0u {
                flagexprs += [make_rt_path_expr(cx, sp, "flag_none")];
            }
            ret make_vec_expr(cx, sp, flagexprs);
        }
        fn make_count(cx: ext_ctxt, sp: span, cnt: count) -> @ast::expr {
            alt cnt {
              count_implied. {
                ret make_rt_path_expr(cx, sp, "count_implied");
              }
              count_is(c) {
                let count_lit = make_new_int(cx, sp, c);
                let count_is_path = make_path_vec(cx, "count_is");
                let count_is_args = [count_lit];
                ret make_call(cx, sp, count_is_path, count_is_args);
              }
              _ { cx.span_unimpl(sp, "unimplemented #fmt conversion"); }
            }
        }
        fn make_ty(cx: ext_ctxt, sp: span, t: ty) -> @ast::expr {
            let rt_type;
            alt t {
              ty_hex(c) {
                alt c {
                  case_upper. { rt_type = "ty_hex_upper"; }
                  case_lower. { rt_type = "ty_hex_lower"; }
                }
              }
              ty_bits. { rt_type = "ty_bits"; }
              ty_octal. { rt_type = "ty_octal"; }
              _ { rt_type = "ty_default"; }
            }
            ret make_rt_path_expr(cx, sp, rt_type);
        }
        fn make_conv_rec(cx: ext_ctxt, sp: span, flags_expr: @ast::expr,
                         width_expr: @ast::expr, precision_expr: @ast::expr,
                         ty_expr: @ast::expr) -> @ast::expr {
            ret make_rec_expr(cx, sp,
                              [{ident: "flags", ex: flags_expr},
                               {ident: "width", ex: width_expr},
                               {ident: "precision", ex: precision_expr},
                               {ident: "ty", ex: ty_expr}]);
        }
        let rt_conv_flags = make_flags(cx, sp, cnv.flags);
        let rt_conv_width = make_count(cx, sp, cnv.width);
        let rt_conv_precision = make_count(cx, sp, cnv.precision);
        let rt_conv_ty = make_ty(cx, sp, cnv.ty);
        ret make_conv_rec(cx, sp, rt_conv_flags, rt_conv_width,
                          rt_conv_precision, rt_conv_ty);
    }
    fn make_conv_call(cx: ext_ctxt, sp: span, conv_type: str, cnv: conv,
                      arg: @ast::expr) -> @ast::expr {
        let fname = "conv_" + conv_type;
        let path = make_path_vec(cx, fname);
        let cnv_expr = make_rt_conv_expr(cx, sp, cnv);
        let args = [cnv_expr, arg];
        ret make_call(cx, arg.span, path, args);
    }
    fn make_new_conv(cx: ext_ctxt, sp: span, cnv: conv, arg: @ast::expr) ->
       @ast::expr {
        // FIXME: Extract all this validation into extfmt::ct

        fn is_signed_type(cnv: conv) -> bool {
            alt cnv.ty {
              ty_int(s) {
                alt s { signed. { ret true; } unsigned. { ret false; } }
              }
              ty_float. { ret true; }
              _ { ret false; }
            }
        }
        let unsupported = "conversion not supported in #fmt string";
        alt cnv.param {
          option::none. { }
          _ { cx.span_unimpl(sp, unsupported); }
        }
        for f: flag in cnv.flags {
            alt f {
              flag_left_justify. { }
              flag_sign_always. {
                if !is_signed_type(cnv) {
                    cx.span_fatal(sp,
                                  "+ flag only valid in " +
                                      "signed #fmt conversion");
                }
              }
              flag_space_for_sign. {
                if !is_signed_type(cnv) {
                    cx.span_fatal(sp,
                                  "space flag only valid in " +
                                      "signed #fmt conversions");
                }
              }
              flag_left_zero_pad. { }
              _ { cx.span_unimpl(sp, unsupported); }
            }
        }
        alt cnv.width {
          count_implied. { }
          count_is(_) { }
          _ { cx.span_unimpl(sp, unsupported); }
        }
        alt cnv.precision {
          count_implied. { }
          count_is(_) { }
          _ { cx.span_unimpl(sp, unsupported); }
        }
        alt cnv.ty {
          ty_str. { ret make_conv_call(cx, arg.span, "str", cnv, arg); }
          ty_int(sign) {
            alt sign {
              signed. { ret make_conv_call(cx, arg.span, "int", cnv, arg); }
              unsigned. {
                ret make_conv_call(cx, arg.span, "uint", cnv, arg);
              }
            }
          }
          ty_bool. { ret make_conv_call(cx, arg.span, "bool", cnv, arg); }
          ty_char. { ret make_conv_call(cx, arg.span, "char", cnv, arg); }
          ty_hex(_) { ret make_conv_call(cx, arg.span, "uint", cnv, arg); }
          ty_bits. { ret make_conv_call(cx, arg.span, "uint", cnv, arg); }
          ty_octal. { ret make_conv_call(cx, arg.span, "uint", cnv, arg); }
          ty_float. { ret make_conv_call(cx, arg.span, "float", cnv, arg); }
          _ { cx.span_unimpl(sp, unsupported); }
        }
    }
    fn log_conv(c: conv) {
        alt c.param {
          some(p) { log(debug, "param: " + int::to_str(p, 10u)); }
          _ { #debug("param: none"); }
        }
        for f: flag in c.flags {
            alt f {
              flag_left_justify. { #debug("flag: left justify"); }
              flag_left_zero_pad. { #debug("flag: left zero pad"); }
              flag_space_for_sign. { #debug("flag: left space pad"); }
              flag_sign_always. { #debug("flag: sign always"); }
              flag_alternate. { #debug("flag: alternate"); }
            }
        }
        alt c.width {
          count_is(i) { log(debug,
                                 "width: count is " + int::to_str(i, 10u)); }
          count_is_param(i) {
            log(debug,
                     "width: count is param " + int::to_str(i, 10u));
          }
          count_is_next_param. { #debug("width: count is next param"); }
          count_implied. { #debug("width: count is implied"); }
        }
        alt c.precision {
          count_is(i) { log(debug,
                                 "prec: count is " + int::to_str(i, 10u)); }
          count_is_param(i) {
            log(debug,
                     "prec: count is param " + int::to_str(i, 10u));
          }
          count_is_next_param. { #debug("prec: count is next param"); }
          count_implied. { #debug("prec: count is implied"); }
        }
        alt c.ty {
          ty_bool. { #debug("type: bool"); }
          ty_str. { #debug("type: str"); }
          ty_char. { #debug("type: char"); }
          ty_int(s) {
            alt s {
              signed. { #debug("type: signed"); }
              unsigned. { #debug("type: unsigned"); }
            }
          }
          ty_bits. { #debug("type: bits"); }
          ty_hex(cs) {
            alt cs {
              case_upper. { #debug("type: uhex"); }
              case_lower. { #debug("type: lhex"); }
            }
          }
          ty_octal. { #debug("type: octal"); }
          ty_float. { #debug("type: float"); }
        }
    }
    let fmt_sp = args[0].span;
    let n = 0u;
    let tmp_expr = make_new_str(cx, sp, "");
    let nargs = vec::len::<@ast::expr>(args);
    for pc: piece in pieces {
        alt pc {
          piece_string(s) {
            let s_expr = make_new_str(cx, fmt_sp, s);
            tmp_expr = make_add_expr(cx, fmt_sp, tmp_expr, s_expr);
          }
          piece_conv(conv) {
            n += 1u;
            if n >= nargs {
                cx.span_fatal(sp,
                              "not enough arguments to #fmt " +
                                  "for the given format string");
            }
            #debug("Building conversion:");
            log_conv(conv);
            let arg_expr = args[n];
            let c_expr = make_new_conv(cx, fmt_sp, conv, arg_expr);
            tmp_expr = make_add_expr(cx, fmt_sp, tmp_expr, c_expr);
          }
        }
    }
    let expected_nargs = n + 1u; // n conversions + the fmt string

    if expected_nargs < nargs {
        cx.span_fatal
            (sp, #fmt["too many arguments to #fmt. found %u, expected %u",
                           nargs, expected_nargs]);
    }
    ret tmp_expr;
}
//
// Local Variables:
// mode: rust
// fill-column: 78;
// indent-tabs-mode: nil
// c-basic-offset: 4
// buffer-file-coding-system: utf-8-unix
// End:
//
<|endoftext|>


<commit_before><commit_msg>Removed a single whitespace<commit_after><|endoftext|>


<commit_before>use std::collections::HashMap;
use std::str::FromStr;

use yaml::Yaml;

pub struct Parameter {
    pub description: Option<String>,
    pub display_name: Option<String>,
    pub name: String,
    pub parameter_type: Option<ParameterType>,
    pub required: bool,
    pub value: Option<String>,
}

pub enum ParameterType {
    Bool,
    Int,
    String,
}

pub type ParamMap = HashMap<String, Parameter>;
pub type UserValues = HashMap<String, String>;

impl Parameter {
    pub fn new(yaml: &Yaml, user_values: &UserValues) -> Result<Self, String> {
        let description = match yaml["description"] {
            Yaml::String(ref description) => Some(description.clone()),
            _ => None,
        };
        let display_name = match yaml["displayName"] {
            Yaml::String(ref description) => Some(description.clone()),
            _ => None,
        };
        let name = match yaml["name"] {
            Yaml::String(ref name) => name.clone(),
            _ => return Err("Parameters must have a \"name\" field.".to_owned()),
        };
        let parameter_type = match yaml["parameterType"].as_str() {
            Some(ref parameter_type) => Some(try!(parameter_type.parse())),
            None => None,
        };
        let required = yaml["required"].as_bool().unwrap_or(false);
        let value = match user_values.get(&name) {
            Some(value) => Some(value.clone()),
            None => match yaml["value"] {
                Yaml::Boolean(ref value)  => Some(format!("{}", value)),
                Yaml::Integer(ref value) => Some(format!("{}", value)),
                Yaml::String(ref value) => Some(value.clone()),
                _ => if required {
                    return Err(
                        format!(
                            "Parameter {} required and must be {}",
                            display_name.unwrap_or(name),
                            parameter_type.map(|pt| match pt {
                                ParameterType::Bool => "a bool",
                                ParameterType::Int => "an int",
                                ParameterType::String => "a string"
                            }).unwrap_or("base64, bool, int, or string")
                        )
                    )
                } else {
                    None
                },
            }
        };

        Ok(Parameter {
            description: description,
            display_name: display_name,
            name: name,
            parameter_type: parameter_type,
            required: required,
            value: value,
        })
    }
}

impl FromStr for ParameterType {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "base64" => Ok(ParameterType::String),
            "bool" => Ok(ParameterType::Bool),
            "int" => Ok(ParameterType::Int),
            "string" => Ok(ParameterType::String),
            _ => Err("parameterType must be base64, bool, int, or string.".to_owned()),
        }
    }
}

pub fn user_values(parameters: Vec<String>) -> Result<UserValues, String> {
    let mut user_values = HashMap::new();

    for parameter in parameters {
        let mut parts: Vec<String> = parameter.split('=').map(|s| s.to_string()).collect();

        if parts.len() < 2 {
            return Err("Parameters must be supplied in the form KEY=VALUE.".to_string());
        } else {
            user_values.insert(parts.remove(0), parts.concat());
        }
    }

    Ok(user_values)
}
<commit_msg>Don't throw away = in parameter values.<commit_after>use std::collections::HashMap;
use std::str::FromStr;

use yaml::Yaml;

pub struct Parameter {
    pub description: Option<String>,
    pub display_name: Option<String>,
    pub name: String,
    pub parameter_type: Option<ParameterType>,
    pub required: bool,
    pub value: Option<String>,
}

pub enum ParameterType {
    Bool,
    Int,
    String,
}

pub type ParamMap = HashMap<String, Parameter>;
pub type UserValues = HashMap<String, String>;

impl Parameter {
    pub fn new(yaml: &Yaml, user_values: &UserValues) -> Result<Self, String> {
        let description = match yaml["description"] {
            Yaml::String(ref description) => Some(description.clone()),
            _ => None,
        };
        let display_name = match yaml["displayName"] {
            Yaml::String(ref description) => Some(description.clone()),
            _ => None,
        };
        let name = match yaml["name"] {
            Yaml::String(ref name) => name.clone(),
            _ => return Err("Parameters must have a \"name\" field.".to_owned()),
        };
        let parameter_type = match yaml["parameterType"].as_str() {
            Some(ref parameter_type) => Some(try!(parameter_type.parse())),
            None => None,
        };
        let required = yaml["required"].as_bool().unwrap_or(false);
        let value = match user_values.get(&name) {
            Some(value) => Some(value.clone()),
            None => match yaml["value"] {
                Yaml::Boolean(ref value)  => Some(format!("{}", value)),
                Yaml::Integer(ref value) => Some(format!("{}", value)),
                Yaml::String(ref value) => Some(value.clone()),
                _ => if required {
                    return Err(
                        format!(
                            "Parameter {} required and must be {}",
                            display_name.unwrap_or(name),
                            parameter_type.map(|pt| match pt {
                                ParameterType::Bool => "a bool",
                                ParameterType::Int => "an int",
                                ParameterType::String => "a string"
                            }).unwrap_or("base64, bool, int, or string")
                        )
                    )
                } else {
                    None
                },
            }
        };

        Ok(Parameter {
            description: description,
            display_name: display_name,
            name: name,
            parameter_type: parameter_type,
            required: required,
            value: value,
        })
    }
}

impl FromStr for ParameterType {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "base64" => Ok(ParameterType::String),
            "bool" => Ok(ParameterType::Bool),
            "int" => Ok(ParameterType::Int),
            "string" => Ok(ParameterType::String),
            _ => Err("parameterType must be base64, bool, int, or string.".to_owned()),
        }
    }
}

pub fn user_values(parameters: Vec<String>) -> Result<UserValues, String> {
    let mut user_values = HashMap::new();

    for parameter in parameters {
        let mut parts: Vec<String> = parameter.split('=').map(|s| s.to_string()).collect();

        if parts.len() < 2 {
            return Err("Parameters must be supplied in the form KEY=VALUE.".to_string());
        } else {
            user_values.insert(parts.remove(0), parts.join("="));
        }
    }

    Ok(user_values)
}
<|endoftext|>


<commit_before>// Copyright 2018 Developers of the Rand project.
// Copyright 2016-2017 The Rust Project Developers.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The Poisson distribution.

use rand::Rng;
use crate::{Distribution, Cauchy, Standard};
use crate::utils::{log_gamma, Float};

/// The Poisson distribution `Poisson(lambda)`.
///
/// This distribution has a density function:
/// `f(k) = lambda^k * exp(-lambda) / k!` for `k >= 0`.
///
/// # Example
///
/// ```
/// use rand_distr::{Poisson, Distribution};
///
/// let poi = Poisson::new(2.0).unwrap();
/// let v: u64 = poi.sample(&mut rand::thread_rng());
/// println!("{} is from a Poisson(2) distribution", v);
/// ```
#[derive(Clone, Copy, Debug)]
pub struct Poisson<N> {
    lambda: N,
    // precalculated values
    exp_lambda: N,
    log_lambda: N,
    sqrt_2lambda: N,
    magic_val: N,
}

/// Error type returned from `Poisson::new`.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
    /// `lambda <= 0` or `nan`.
    ShapeTooSmall,
}

impl<N: Float> Poisson<N> {
    /// Construct a new `Poisson` with the given shape parameter
    /// `lambda`.
    pub fn new(lambda: N) -> Result<Poisson<N>, Error> {
        if !(lambda > N::from(0.0)) {
            return Err(Error::ShapeTooSmall);
        }
        let log_lambda = lambda.ln();
        Ok(Poisson {
            lambda,
            exp_lambda: (-lambda).exp(),
            log_lambda,
            sqrt_2lambda: (N::from(2.0) * lambda).sqrt(),
            magic_val: lambda * log_lambda - log_gamma(N::from(1.0) + lambda),
        })
    }
}

impl<N: Float> Distribution<N> for Poisson<N>
where Standard: Distribution<N>
{
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> N {
        // using the algorithm from Numerical Recipes in C

        // for low expected values use the Knuth method
        if self.lambda < N::from(12.0) {
            let mut result = N::from(0.);
            let mut p = N::from(1.0);
            while p > self.exp_lambda {
                p *= rng.gen::<N>();
                result += N::from(1.);
            }
            result - N::from(1.)
        }
        // high expected values - rejection method
        else {
            // we use the Cauchy distribution as the comparison distribution
            // f(x) ~ 1/(1+x^2)
            let cauchy = Cauchy::new(N::from(0.0), N::from(1.0)).unwrap();
            let mut result;

            loop {
                let mut comp_dev;

                loop {
                    // draw from the Cauchy distribution
                    comp_dev = rng.sample(cauchy);
                    // shift the peak of the comparison ditribution
                    result = self.sqrt_2lambda * comp_dev + self.lambda;
                    // repeat the drawing until we are in the range of possible values
                    if result >= N::from(0.0) {
                        break;
                    }
                }
                // now the result is a random variable greater than 0 with Cauchy distribution
                // the result should be an integer value
                result = result.floor();

                // this is the ratio of the Poisson distribution to the comparison distribution
                // the magic value scales the distribution function to a range of approximately 0-1
                // since it is not exact, we multiply the ratio by 0.9 to avoid ratios greater than 1
                // this doesn't change the resulting distribution, only increases the rate of failed drawings
                let check = N::from(0.9) * (N::from(1.0) + comp_dev * comp_dev)
                    * (result * self.log_lambda - log_gamma(N::from(1.0) + result) - self.magic_val).exp();

                // check with uniform random value - if below the threshold, we are within the target distribution
                if rng.gen::<N>() <= check {
                    break;
                }
            }
            result
        }
    }
}

impl<N: Float> Distribution<u64> for Poisson<N>
where Standard: Distribution<N>
{
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u64 {
        let result: N = self.sample(rng);
        result.into_ui().unwrap()
    }
}

#[cfg(test)]
mod test {
    use crate::Distribution;
    use super::Poisson;

    #[test]
    fn test_poisson_10() {
        let poisson = Poisson::new(10.0).unwrap();
        let mut rng = crate::test::rng(123);
        let mut sum = 0;
        for _ in 0..1000 {
            let s: u64 = poisson.sample(&mut rng);
            sum += s;
        }
        let avg = (sum as f64) / 1000.0;
        println!("Poisson average: {}", avg);
        assert!((avg - 10.0).abs() < 0.5); // not 100% certain, but probable enough
    }

    #[test]
    fn test_poisson_15() {
        // Take the 'high expected values' path
        let poisson = Poisson::new(15.0).unwrap();
        let mut rng = crate::test::rng(123);
        let mut sum = 0;
        for _ in 0..1000 {
            let s: u64 = poisson.sample(&mut rng);
            sum += s;
        }
        let avg = (sum as f64) / 1000.0;
        println!("Poisson average: {}", avg);
        assert!((avg - 15.0).abs() < 0.5); // not 100% certain, but probable enough
    }

    #[test]
    #[should_panic]
    fn test_poisson_invalid_lambda_zero() {
        Poisson::new(0.0).unwrap();
    }

    #[test]
    #[should_panic]
    fn test_poisson_invalid_lambda_neg() {
        Poisson::new(-10.0).unwrap();
    }
}
<commit_msg>rand_distr: Add tests for f64 sampling from Poisson<commit_after>// Copyright 2018 Developers of the Rand project.
// Copyright 2016-2017 The Rust Project Developers.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The Poisson distribution.

use rand::Rng;
use crate::{Distribution, Cauchy, Standard};
use crate::utils::{log_gamma, Float};

/// The Poisson distribution `Poisson(lambda)`.
///
/// This distribution has a density function:
/// `f(k) = lambda^k * exp(-lambda) / k!` for `k >= 0`.
///
/// # Example
///
/// ```
/// use rand_distr::{Poisson, Distribution};
///
/// let poi = Poisson::new(2.0).unwrap();
/// let v: u64 = poi.sample(&mut rand::thread_rng());
/// println!("{} is from a Poisson(2) distribution", v);
/// ```
#[derive(Clone, Copy, Debug)]
pub struct Poisson<N> {
    lambda: N,
    // precalculated values
    exp_lambda: N,
    log_lambda: N,
    sqrt_2lambda: N,
    magic_val: N,
}

/// Error type returned from `Poisson::new`.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
    /// `lambda <= 0` or `nan`.
    ShapeTooSmall,
}

impl<N: Float> Poisson<N> {
    /// Construct a new `Poisson` with the given shape parameter
    /// `lambda`.
    pub fn new(lambda: N) -> Result<Poisson<N>, Error> {
        if !(lambda > N::from(0.0)) {
            return Err(Error::ShapeTooSmall);
        }
        let log_lambda = lambda.ln();
        Ok(Poisson {
            lambda,
            exp_lambda: (-lambda).exp(),
            log_lambda,
            sqrt_2lambda: (N::from(2.0) * lambda).sqrt(),
            magic_val: lambda * log_lambda - log_gamma(N::from(1.0) + lambda),
        })
    }
}

impl<N: Float> Distribution<N> for Poisson<N>
where Standard: Distribution<N>
{
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> N {
        // using the algorithm from Numerical Recipes in C

        // for low expected values use the Knuth method
        if self.lambda < N::from(12.0) {
            let mut result = N::from(0.);
            let mut p = N::from(1.0);
            while p > self.exp_lambda {
                p *= rng.gen::<N>();
                result += N::from(1.);
            }
            result - N::from(1.)
        }
        // high expected values - rejection method
        else {
            // we use the Cauchy distribution as the comparison distribution
            // f(x) ~ 1/(1+x^2)
            let cauchy = Cauchy::new(N::from(0.0), N::from(1.0)).unwrap();
            let mut result;

            loop {
                let mut comp_dev;

                loop {
                    // draw from the Cauchy distribution
                    comp_dev = rng.sample(cauchy);
                    // shift the peak of the comparison ditribution
                    result = self.sqrt_2lambda * comp_dev + self.lambda;
                    // repeat the drawing until we are in the range of possible values
                    if result >= N::from(0.0) {
                        break;
                    }
                }
                // now the result is a random variable greater than 0 with Cauchy distribution
                // the result should be an integer value
                result = result.floor();

                // this is the ratio of the Poisson distribution to the comparison distribution
                // the magic value scales the distribution function to a range of approximately 0-1
                // since it is not exact, we multiply the ratio by 0.9 to avoid ratios greater than 1
                // this doesn't change the resulting distribution, only increases the rate of failed drawings
                let check = N::from(0.9) * (N::from(1.0) + comp_dev * comp_dev)
                    * (result * self.log_lambda - log_gamma(N::from(1.0) + result) - self.magic_val).exp();

                // check with uniform random value - if below the threshold, we are within the target distribution
                if rng.gen::<N>() <= check {
                    break;
                }
            }
            result
        }
    }
}

impl<N: Float> Distribution<u64> for Poisson<N>
where Standard: Distribution<N>
{
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u64 {
        let result: N = self.sample(rng);
        result.into_ui().unwrap()
    }
}

#[cfg(test)]
mod test {
    use crate::Distribution;
    use super::Poisson;

    #[test]
    fn test_poisson_10() {
        let poisson = Poisson::new(10.0).unwrap();
        let mut rng = crate::test::rng(123);
        let mut sum_u64 = 0;
        let mut sum_f64 = 0.;
        for _ in 0..1000 {
            let s_u64: u64 = poisson.sample(&mut rng);
            let s_f64: f64 = poisson.sample(&mut rng);
            sum_u64 += s_u64;
            sum_f64 += s_f64;
        }
        let avg_u64 = (sum_u64 as f64) / 1000.0;
        let avg_f64 = sum_f64 / 1000.0;
        println!("Poisson averages: {} (u64)  {} (f64)", avg_u64, avg_f64);
        for &avg in &[avg_u64, avg_f64] {
            assert!((avg - 10.0).abs() < 0.5); // not 100% certain, but probable enough
        }
    }

    #[test]
    fn test_poisson_15() {
        // Take the 'high expected values' path
        let poisson = Poisson::new(15.0).unwrap();
        let mut rng = crate::test::rng(123);
        let mut sum_u64 = 0;
        let mut sum_f64 = 0.;
        for _ in 0..1000 {
            let s_u64: u64 = poisson.sample(&mut rng);
            let s_f64: f64 = poisson.sample(&mut rng);
            sum_u64 += s_u64;
            sum_f64 += s_f64;
        }
        let avg_u64 = (sum_u64 as f64) / 1000.0;
        let avg_f64 = sum_f64 / 1000.0;
        println!("Poisson average: {} (u64)  {} (f64)", avg_u64, avg_f64);
        for &avg in &[avg_u64, avg_f64] {
            assert!((avg - 15.0).abs() < 0.5); // not 100% certain, but probable enough
        }
    }

    #[test]
    #[should_panic]
    fn test_poisson_invalid_lambda_zero() {
        Poisson::new(0.0).unwrap();
    }

    #[test]
    #[should_panic]
    fn test_poisson_invalid_lambda_neg() {
        Poisson::new(-10.0).unwrap();
    }
}
<|endoftext|>