<commit_before>//
// Copyright 2018 The Project Oak Authors
//
// 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.
//

//! Functionality to help Oak Nodes interact with gRPC.

pub use crate::proto::code::Code;
use crate::{proto, Handle, OakError, OakStatus, ReadHandle, WriteHandle};
use log::info;
use protobuf::{Message, ProtobufEnum};

/// Result type that uses a [`proto::status::Status`] type for error values.
pub type Result<T> = std::result::Result<T, proto::status::Status>;

/// Helper to create a gRPC status object.
pub fn build_status(code: Code, msg: &str) -> proto::status::Status {
    let mut status = proto::status::Status::new();
    status.set_code(code.value());
    status.set_message(msg.to_string());
    status
}

/// Channel-holding object that encapsulates response messages into
/// `GrpcResponse` wrapper messages and writes serialized versions to a send
///  channel.
pub struct ChannelResponseWriter {
    channel: crate::io::Sender,
}

/// Indicate whether a write method should leave the current gRPC method
/// invocation open or close it.
#[derive(PartialEq, Clone, Debug)]
pub enum WriteMode {
    KeepOpen,
    Close,
}

impl ChannelResponseWriter {
    pub fn new(out_handle: crate::WriteHandle) -> Self {
        ChannelResponseWriter {
            channel: crate::io::Sender::new(out_handle),
        }
    }

    /// Retrieve the Oak handle underlying the writer object.
    pub fn handle(self) -> crate::WriteHandle {
        self.channel.handle
    }

    /// Write out a gRPC response and optionally close out the method
    /// invocation.  Any errors from the channel are silently dropped.
    pub fn write<T: protobuf::Message>(
        &mut self,
        rsp: &T,
        mode: WriteMode,
    ) -> std::result::Result<(), OakError> {
        // Put the serialized response into a GrpcResponse message wrapper and
        // serialize it into the channel.
        let mut grpc_rsp = proto::grpc_encap::GrpcResponse::new();
        let mut any = protobuf::well_known_types::Any::new();
        rsp.write_to_writer(&mut any.value)?;
        grpc_rsp.set_rsp_msg(any);
        grpc_rsp.set_last(match mode {
            WriteMode::KeepOpen => false,
            WriteMode::Close => true,
        });
        grpc_rsp.write_to_writer(&mut self.channel)?;
        if mode == WriteMode::Close {
            self.channel.close()?;
        }
        Ok(())
    }

    /// Write an empty gRPC response and optionally close out the method
    /// invocation. Any errors from the channel are silently dropped.
    pub fn write_empty(&mut self, mode: WriteMode) -> std::result::Result<(), OakError> {
        let mut grpc_rsp = proto::grpc_encap::GrpcResponse::new();
        grpc_rsp.set_rsp_msg(protobuf::well_known_types::Any::new());
        grpc_rsp.set_last(match mode {
            WriteMode::KeepOpen => false,
            WriteMode::Close => true,
        });
        grpc_rsp.write_to_writer(&mut self.channel)?;
        if mode == WriteMode::Close {
            self.channel.close()?;
        }
        Ok(())
    }

    /// Close out the gRPC method invocation with the given final result.
    pub fn close(&mut self, result: Result<()>) -> std::result::Result<(), OakError> {
        // Build a final GrpcResponse message wrapper and serialize it into the
        // channel.
        let mut grpc_rsp = proto::grpc_encap::GrpcResponse::new();
        grpc_rsp.set_last(true);
        if let Err(status) = result {
            grpc_rsp.set_status(status);
        }
        grpc_rsp.write_to_writer(&mut self.channel)?;
        self.channel.close()?;
        Ok(())
    }
}

/// Trait for Oak Nodes that act as a gRPC services.
///
/// An `OakNode` instance is normally passed to [`event_loop`], to allow
/// repeated invocation of its `invoke()` method.
pub trait OakNode {
    /// Construct the (single) instance of the node.
    ///
    /// This method may choose to initialize logging by invoking
    /// [`oak_log::init()`].
    ///
    /// [`oak_log::init()`]: ../../oak_log/fn.init.html
    fn new() -> Self
    where
        Self: Sized;

    /// Process a single gRPC method invocation.
    ///
    /// The method name is provided by `method` and the incoming serialized gRPC
    /// request is held in `req`.  Response messages should be written using
    /// `writer.write`, followed by `writer.close`.
    fn invoke(&mut self, method: &str, req: &[u8], writer: ChannelResponseWriter);
}

/// Perform a gRPC event loop for a Node.
///
/// Invoking the given `node`'s [`invoke`] method for each incoming request that
/// arrives on the inbound channel as a serialized [`GrpcRequest`] message,
/// giving the [`invoke`] method the outbound channel for encapsulated responses
/// to be written to.
///
/// [`invoke`]: OakNode::invoke
/// [`GrpcRequest`]: crate::proto::grpc_encap::GrpcRequest
pub fn event_loop<T: OakNode>(
    mut node: T,
    grpc_in_handle: ReadHandle,
) -> std::result::Result<(), crate::OakStatus> {
    info!("start event loop for node with handle {:?}", grpc_in_handle);
    if !grpc_in_handle.handle.is_valid() {
        return Err(OakStatus::ERR_CHANNEL_CLOSED);
    }
    crate::set_panic_hook();

    let read_handles = vec![grpc_in_handle];
    let mut space = crate::new_handle_space(&read_handles);

    loop {
        // Block until there is a message to read on an input channel.
        crate::prep_handle_space(&mut space);
        // TODO: Use higher-level wait function from SDK instead of the ABI one.
        let status =
            unsafe { oak_abi::wait_on_channels(space.as_mut_ptr(), read_handles.len() as u32) };
        crate::result_from_status(status as i32, ())?;

        let mut buf = Vec::<u8>::with_capacity(1024);
        let mut handles = Vec::<Handle>::with_capacity(1);
        crate::channel_read(grpc_in_handle, &mut buf, &mut handles)?;
        if buf.is_empty() {
            panic!("no bytes received")
        }
        if handles.is_empty() {
            panic!("no response handle received alongside gRPC request")
        }
        let req: proto::grpc_encap::GrpcRequest =
            protobuf::parse_from_bytes(&buf).expect("failed to parse GrpcRequest message");
        if !req.last {
            panic!("Support for streaming requests not yet implemented");
        }
        node.invoke(
            &req.method_name,
            req.get_req_msg().value.as_slice(),
            ChannelResponseWriter::new(WriteHandle { handle: handles[0] }),
        );
    }
}

/// Generic function that handles request deserialization and response
/// serialization (and sending down a channel) for a function that handles a
/// request/response pair.
///
/// Panics if [de-]serialization or channel operations fail.
pub fn handle_req_rsp<C, R, Q>(mut node_fn: C, req: &[u8], mut writer: ChannelResponseWriter)
where
    C: FnMut(R) -> Result<Q>,
    R: protobuf::Message,
    Q: protobuf::Message,
{
    let r: R = protobuf::parse_from_bytes(&req).expect("Failed to parse request protobuf message");
    let result = match node_fn(r) {
        Ok(rsp) => writer.write(&rsp, WriteMode::Close),
        Err(status) => writer.close(Err(status)),
    };
    if let Err(e) = result {
        panic!("Failed to process response: {:?}", e)
    }
}

/// Generic function that handles request deserialization and response
/// serialization (and sending down a channel) for a function that handles a
/// request and streams responses.
///
/// Panics if [de-]serialization or channel operations fail.
pub fn handle_req_stream<C, R>(mut node_fn: C, req: &[u8], writer: ChannelResponseWriter)
where
    C: FnMut(R, ChannelResponseWriter),
    R: protobuf::Message,
{
    let r: R = protobuf::parse_from_bytes(&req).expect("Failed to parse request protobuf message");
    node_fn(r, writer)
}

/// Generic function that handles request deserialization and response
/// serialization (and sending down a channel) for a function that handles a
/// stream of requests to produce a response.
///
/// Panics if [de-]serialization or channel operations fail.
pub fn handle_stream_rsp<C, R, Q>(mut node_fn: C, req: &[u8], mut writer: ChannelResponseWriter)
where
    C: FnMut(Vec<R>) -> Result<Q>,
    R: protobuf::Message,
    Q: protobuf::Message,
{
    // TODO(#97): better client-side streaming
    let rr: Vec<R> =
        vec![protobuf::parse_from_bytes(&req).expect("Failed to parse request protobuf message")];
    let result = match node_fn(rr) {
        Ok(rsp) => writer.write(&rsp, WriteMode::Close),
        Err(status) => writer.close(Err(status)),
    };
    if let Err(e) = result {
        panic!("Failed to process response: {:?}", e)
    }
}

/// Generic function that handles request deserialization and response
/// serialization (and sending down a channel) for a function that handles a
/// stream of requests to produce a stream of responses.
///
/// Panics if [de-]serialization or channel operations fail.
pub fn handle_stream_stream<C, R>(mut node_fn: C, req: &[u8], writer: ChannelResponseWriter)
where
    C: FnMut(Vec<R>, ChannelResponseWriter),
    R: protobuf::Message,
{
    // TODO(#97): better client-side streaming
    let rr: Vec<R> =
        vec![protobuf::parse_from_bytes(&req).expect("Failed to parse request protobuf message")];
    node_fn(rr, writer)
}
<commit_msg>sdk: don't set panic hook in grpc helper<commit_after>//
// Copyright 2018 The Project Oak Authors
//
// 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.
//

//! Functionality to help Oak Nodes interact with gRPC.

pub use crate::proto::code::Code;
use crate::{proto, Handle, OakError, OakStatus, ReadHandle, WriteHandle};
use log::info;
use protobuf::{Message, ProtobufEnum};

/// Result type that uses a [`proto::status::Status`] type for error values.
pub type Result<T> = std::result::Result<T, proto::status::Status>;

/// Helper to create a gRPC status object.
pub fn build_status(code: Code, msg: &str) -> proto::status::Status {
    let mut status = proto::status::Status::new();
    status.set_code(code.value());
    status.set_message(msg.to_string());
    status
}

/// Channel-holding object that encapsulates response messages into
/// `GrpcResponse` wrapper messages and writes serialized versions to a send
///  channel.
pub struct ChannelResponseWriter {
    channel: crate::io::Sender,
}

/// Indicate whether a write method should leave the current gRPC method
/// invocation open or close it.
#[derive(PartialEq, Clone, Debug)]
pub enum WriteMode {
    KeepOpen,
    Close,
}

impl ChannelResponseWriter {
    pub fn new(out_handle: crate::WriteHandle) -> Self {
        ChannelResponseWriter {
            channel: crate::io::Sender::new(out_handle),
        }
    }

    /// Retrieve the Oak handle underlying the writer object.
    pub fn handle(self) -> crate::WriteHandle {
        self.channel.handle
    }

    /// Write out a gRPC response and optionally close out the method
    /// invocation.  Any errors from the channel are silently dropped.
    pub fn write<T: protobuf::Message>(
        &mut self,
        rsp: &T,
        mode: WriteMode,
    ) -> std::result::Result<(), OakError> {
        // Put the serialized response into a GrpcResponse message wrapper and
        // serialize it into the channel.
        let mut grpc_rsp = proto::grpc_encap::GrpcResponse::new();
        let mut any = protobuf::well_known_types::Any::new();
        rsp.write_to_writer(&mut any.value)?;
        grpc_rsp.set_rsp_msg(any);
        grpc_rsp.set_last(match mode {
            WriteMode::KeepOpen => false,
            WriteMode::Close => true,
        });
        grpc_rsp.write_to_writer(&mut self.channel)?;
        if mode == WriteMode::Close {
            self.channel.close()?;
        }
        Ok(())
    }

    /// Write an empty gRPC response and optionally close out the method
    /// invocation. Any errors from the channel are silently dropped.
    pub fn write_empty(&mut self, mode: WriteMode) -> std::result::Result<(), OakError> {
        let mut grpc_rsp = proto::grpc_encap::GrpcResponse::new();
        grpc_rsp.set_rsp_msg(protobuf::well_known_types::Any::new());
        grpc_rsp.set_last(match mode {
            WriteMode::KeepOpen => false,
            WriteMode::Close => true,
        });
        grpc_rsp.write_to_writer(&mut self.channel)?;
        if mode == WriteMode::Close {
            self.channel.close()?;
        }
        Ok(())
    }

    /// Close out the gRPC method invocation with the given final result.
    pub fn close(&mut self, result: Result<()>) -> std::result::Result<(), OakError> {
        // Build a final GrpcResponse message wrapper and serialize it into the
        // channel.
        let mut grpc_rsp = proto::grpc_encap::GrpcResponse::new();
        grpc_rsp.set_last(true);
        if let Err(status) = result {
            grpc_rsp.set_status(status);
        }
        grpc_rsp.write_to_writer(&mut self.channel)?;
        self.channel.close()?;
        Ok(())
    }
}

/// Trait for Oak Nodes that act as a gRPC services.
///
/// An `OakNode` instance is normally passed to [`event_loop`], to allow
/// repeated invocation of its `invoke()` method.
pub trait OakNode {
    /// Construct the (single) instance of the node.
    ///
    /// This method may choose to initialize logging by invoking
    /// [`oak_log::init()`].
    ///
    /// [`oak_log::init()`]: ../../oak_log/fn.init.html
    fn new() -> Self
    where
        Self: Sized;

    /// Process a single gRPC method invocation.
    ///
    /// The method name is provided by `method` and the incoming serialized gRPC
    /// request is held in `req`.  Response messages should be written using
    /// `writer.write`, followed by `writer.close`.
    fn invoke(&mut self, method: &str, req: &[u8], writer: ChannelResponseWriter);
}

/// Perform a gRPC event loop for a Node.
///
/// Invoking the given `node`'s [`invoke`] method for each incoming request that
/// arrives on the inbound channel as a serialized [`GrpcRequest`] message,
/// giving the [`invoke`] method the outbound channel for encapsulated responses
/// to be written to.
///
/// [`invoke`]: OakNode::invoke
/// [`GrpcRequest`]: crate::proto::grpc_encap::GrpcRequest
pub fn event_loop<T: OakNode>(
    mut node: T,
    grpc_in_handle: ReadHandle,
) -> std::result::Result<(), crate::OakStatus> {
    info!("start event loop for node with handle {:?}", grpc_in_handle);
    if !grpc_in_handle.handle.is_valid() {
        return Err(OakStatus::ERR_CHANNEL_CLOSED);
    }
    let read_handles = vec![grpc_in_handle];
    let mut space = crate::new_handle_space(&read_handles);

    loop {
        // Block until there is a message to read on an input channel.
        crate::prep_handle_space(&mut space);
        // TODO: Use higher-level wait function from SDK instead of the ABI one.
        let status =
            unsafe { oak_abi::wait_on_channels(space.as_mut_ptr(), read_handles.len() as u32) };
        crate::result_from_status(status as i32, ())?;

        let mut buf = Vec::<u8>::with_capacity(1024);
        let mut handles = Vec::<Handle>::with_capacity(1);
        crate::channel_read(grpc_in_handle, &mut buf, &mut handles)?;
        if buf.is_empty() {
            panic!("no bytes received")
        }
        if handles.is_empty() {
            panic!("no response handle received alongside gRPC request")
        }
        let req: proto::grpc_encap::GrpcRequest =
            protobuf::parse_from_bytes(&buf).expect("failed to parse GrpcRequest message");
        if !req.last {
            panic!("Support for streaming requests not yet implemented");
        }
        node.invoke(
            &req.method_name,
            req.get_req_msg().value.as_slice(),
            ChannelResponseWriter::new(WriteHandle { handle: handles[0] }),
        );
    }
}

/// Generic function that handles request deserialization and response
/// serialization (and sending down a channel) for a function that handles a
/// request/response pair.
///
/// Panics if [de-]serialization or channel operations fail.
pub fn handle_req_rsp<C, R, Q>(mut node_fn: C, req: &[u8], mut writer: ChannelResponseWriter)
where
    C: FnMut(R) -> Result<Q>,
    R: protobuf::Message,
    Q: protobuf::Message,
{
    let r: R = protobuf::parse_from_bytes(&req).expect("Failed to parse request protobuf message");
    let result = match node_fn(r) {
        Ok(rsp) => writer.write(&rsp, WriteMode::Close),
        Err(status) => writer.close(Err(status)),
    };
    if let Err(e) = result {
        panic!("Failed to process response: {:?}", e)
    }
}

/// Generic function that handles request deserialization and response
/// serialization (and sending down a channel) for a function that handles a
/// request and streams responses.
///
/// Panics if [de-]serialization or channel operations fail.
pub fn handle_req_stream<C, R>(mut node_fn: C, req: &[u8], writer: ChannelResponseWriter)
where
    C: FnMut(R, ChannelResponseWriter),
    R: protobuf::Message,
{
    let r: R = protobuf::parse_from_bytes(&req).expect("Failed to parse request protobuf message");
    node_fn(r, writer)
}

/// Generic function that handles request deserialization and response
/// serialization (and sending down a channel) for a function that handles a
/// stream of requests to produce a response.
///
/// Panics if [de-]serialization or channel operations fail.
pub fn handle_stream_rsp<C, R, Q>(mut node_fn: C, req: &[u8], mut writer: ChannelResponseWriter)
where
    C: FnMut(Vec<R>) -> Result<Q>,
    R: protobuf::Message,
    Q: protobuf::Message,
{
    // TODO(#97): better client-side streaming
    let rr: Vec<R> =
        vec![protobuf::parse_from_bytes(&req).expect("Failed to parse request protobuf message")];
    let result = match node_fn(rr) {
        Ok(rsp) => writer.write(&rsp, WriteMode::Close),
        Err(status) => writer.close(Err(status)),
    };
    if let Err(e) = result {
        panic!("Failed to process response: {:?}", e)
    }
}

/// Generic function that handles request deserialization and response
/// serialization (and sending down a channel) for a function that handles a
/// stream of requests to produce a stream of responses.
///
/// Panics if [de-]serialization or channel operations fail.
pub fn handle_stream_stream<C, R>(mut node_fn: C, req: &[u8], writer: ChannelResponseWriter)
where
    C: FnMut(Vec<R>, ChannelResponseWriter),
    R: protobuf::Message,
{
    // TODO(#97): better client-side streaming
    let rr: Vec<R> =
        vec![protobuf::parse_from_bytes(&req).expect("Failed to parse request protobuf message")];
    node_fn(rr, writer)
}
<|endoftext|>


<commit_before>use std::f64;
use rand::Rng;
use rand::distributions::{Sample, IndependentSample};
use error::StatsError;
use function::{beta, gamma};
use result::Result;
use super::*;

/// Implements the [Beta](https://en.wikipedia.org/wiki/Beta_distribution) distribution
///
/// # Examples
///
/// ```
/// use statrs::distribution::{Beta};
///
/// let n = Beta::new(2.0, 2.0);
/// ```
#[derive(Debug, Copy, Clone, PartialEq)]
pub struct Beta {
    shape_a: f64,
    shape_b: f64,
}

impl Beta {
    /// Constructs a new beta distribution with shapeA (α) of `shape_a`
    /// and shapeB (β) of `shape_b`
    ///
    /// # Errors
    ///
    /// Returns an error if `shape_a` or `shape_b` are `NaN`.
    /// Also returns an error if `shape_a <= 0.0` or `shape_b <= 0.0`
    ///
    /// # Examples
    ///
    /// ```
    /// use statrs::distribution::Beta;
    ///
    /// let mut result = Beta::new(2.0, 2.0);
    /// assert!(result.is_ok());
    ///
    /// result = Beta::new(0.0, 0.0);
    /// assert!(result.is_err());
    /// ```
    pub fn new(shape_a: f64, shape_b: f64) -> Result<Beta> {
        let is_nan = shape_a.is_nan() || shape_b.is_nan();
        match (shape_a, shape_b, is_nan) {
            (_, _, true) => Err(StatsError::BadParams),
            (_, _, false) if shape_a <= 0.0 || shape_b <= 0.0 => Err(StatsError::BadParams),
            (_, _, false) => {
                Ok(Beta {
                    shape_a: shape_a,
                    shape_b: shape_b,
                })
            }
        }
    }

    /// Returns the shapeA (α) of the beta distribution
    ///
    /// # Examples
    ///
    /// ```
    /// use statrs::distribution::Beta;
    ///
    /// let n = Beta::new(2.0, 2.0).unwrap();
    /// assert_eq!(n.shape_a(), 2.0);
    /// ```
    pub fn shape_a(&self) -> f64 {
        self.shape_a
    }

    /// Returns the shapeB (β) of the beta distributionβ
    ///
    /// # Examples
    ///
    /// ```
    /// use statrs::distribution::Beta;
    ///
    /// let n = Beta::new(2.0, 2.0).unwrap();
    /// assert_eq!(n.shape_b(), 2.0);
    /// ```
    pub fn shape_b(&self) -> f64 {
        self.shape_b
    }
}

impl Sample<f64> for Beta {
    /// Generate a random sample from a beta distribution
    /// using `r` as the source of randomness.
    /// Refer [here](#method.sample-1) for implementation details.
    fn sample<R: Rng>(&mut self, r: &mut R) -> f64 {
        super::Distribution::sample(self, r)
    }
}

impl IndependentSample<f64> for Beta {
    /// Generate a random independent sample from a beta distribution
    /// using `r` as the source of randomness.
    /// Refer [here](#method.sample-1) for implementation details.
    fn ind_sample<R: Rng>(&self, r: &mut R) -> f64 {
        super::Distribution::sample(self, r)
    }
}

impl Distribution<f64> for Beta {
    /// Generate a random sample from a beta distribution using
    /// `r` as the source of randomness. Generated by sampling
    /// two gamma distributions and normalizing.
    ///
    /// # Examples
    ///
    /// ```
    /// # extern crate rand;
    /// # extern crate statrs;
    /// use rand::StdRng;
    /// use statrs::distribution::{Beta, Distribution};
    ///
    /// # fn main() {
    /// let mut r = rand::StdRng::new().unwrap();
    /// let n = Beta::new(2.0, 2.0).unwrap();
    /// print!("{}", n.sample::<StdRng>(&mut r));
    /// # }
    /// ```
    fn sample<R: Rng>(&self, r: &mut R) -> f64 {
        let x = super::gamma::sample_unchecked(r, self.shape_a, 1.0);
        let y = super::gamma::sample_unchecked(r, self.shape_b, 1.0);
        x / (x + y)
    }
}

impl Univariate<f64, f64> for Beta {
    /// Calculates the cumulative distribution function for the beta distribution
    /// at `x`
    ///
    /// # Panics
    ///
    /// If `x < 0.0` or `x > 1.0`
    ///
    /// # Formula
    ///
    /// ```ignore
    /// I_x(α, β)
    /// ```
    ///
    /// where `α` is shapeA, `β` is shapeB, and `I_x` is the regularized
    /// lower incomplete beta function
    fn cdf(&self, x: f64) -> f64 {
        assert!(x >= 0.0 && x <= 1.0,
                format!("{}", StatsError::ArgIntervalIncl("x", 0.0, 1.0)));
        if x == 1.0 {
            1.0
        } else if self.shape_a == f64::INFINITY && self.shape_b == f64::INFINITY {
            if x < 0.5 { 0.0 } else { 1.0 }
        } else if self.shape_a == f64::INFINITY {
            if x < 1.0 { 0.0 } else { 1.0 }
        } else if self.shape_b == f64::INFINITY {
            0.0
        } else if self.shape_a == 1.0 && self.shape_b == 1.0 {
            x
        } else {
            beta::beta_reg(self.shape_a, self.shape_b, x)
        }
    }

    /// Returns the minimum value in the domain of the
    /// beta distribution representable by a double precision
    /// float
    ///
    /// # Formula
    ///
    /// ```ignore
    /// 0
    /// ```
    fn min(&self) -> f64 {
        0.0
    }

    /// Returns the maximum value in the domain of the
    /// beta distribution representable by a double precision
    /// float
    ///
    /// # Formula
    ///
    /// ```ignore
    /// 1
    /// ```
    fn max(&self) -> f64 {
        1.0
    }
}

impl Mean<f64, f64> for Beta {
    /// Returns the mean of the beta distribution
    ///
    /// # Formula
    ///
    /// ```ignore
    /// α / (α + β)
    /// ```
    ///
    /// where `α` is shapeA and `β` is shapeB
    fn mean(&self) -> f64 {
        if self.shape_a == f64::INFINITY && self.shape_b == f64::INFINITY {
            0.5
        } else if self.shape_a == f64::INFINITY {
            1.0
        } else if self.shape_b == f64::INFINITY {
            0.0
        } else {
            self.shape_a / (self.shape_a + self.shape_b)
        }
    }
}

impl Variance<f64, f64> for Beta {
    /// Returns the variance of the beta distribution
    ///
    /// # Formula
    ///
    /// ```ignore
    /// (α * β) / ((α + β)^2 * (α + β + 1))
    /// ```
    ///
    /// where `α` is shapeA and `β` is shapeB
    fn variance(&self) -> f64 {
        self.shape_a * self.shape_b /
        ((self.shape_a + self.shape_b) * (self.shape_a + self.shape_b) *
         (self.shape_a + self.shape_b + 1.0))
    }
}

impl Entropy<f64> for Beta {
    /// Returns the entropy of the beta distribution
    ///
    /// # Formula
    ///
    /// ```ignore
    /// ln(B(α, β)) - (α - 1)ψ(α) - (β - 1)ψ(β) + (α + β - 2)ψ(α + β)
    /// ```
    ///
    /// where `α` is shapeA, `β` is shapeB and `ψ` is the digamma function
    fn entropy(&self) -> f64 {
        if self.shape_a == f64::INFINITY || self.shape_b == f64::INFINITY {
            0.0
        } else {
            beta::ln_beta(self.shape_a, self.shape_b) -
            (self.shape_a - 1.0) * gamma::digamma(self.shape_a) -
            (self.shape_b - 1.0) * gamma::digamma(self.shape_b) +
            (self.shape_a + self.shape_b - 2.0) * gamma::digamma(self.shape_a + self.shape_b)
        }
    }
}
<commit_msg>forgot std dev for variance impl<commit_after>use std::f64;
use rand::Rng;
use rand::distributions::{Sample, IndependentSample};
use error::StatsError;
use function::{beta, gamma};
use result::Result;
use super::*;

/// Implements the [Beta](https://en.wikipedia.org/wiki/Beta_distribution) distribution
///
/// # Examples
///
/// ```
/// use statrs::distribution::{Beta};
///
/// let n = Beta::new(2.0, 2.0);
/// ```
#[derive(Debug, Copy, Clone, PartialEq)]
pub struct Beta {
    shape_a: f64,
    shape_b: f64,
}

impl Beta {
    /// Constructs a new beta distribution with shapeA (α) of `shape_a`
    /// and shapeB (β) of `shape_b`
    ///
    /// # Errors
    ///
    /// Returns an error if `shape_a` or `shape_b` are `NaN`.
    /// Also returns an error if `shape_a <= 0.0` or `shape_b <= 0.0`
    ///
    /// # Examples
    ///
    /// ```
    /// use statrs::distribution::Beta;
    ///
    /// let mut result = Beta::new(2.0, 2.0);
    /// assert!(result.is_ok());
    ///
    /// result = Beta::new(0.0, 0.0);
    /// assert!(result.is_err());
    /// ```
    pub fn new(shape_a: f64, shape_b: f64) -> Result<Beta> {
        let is_nan = shape_a.is_nan() || shape_b.is_nan();
        match (shape_a, shape_b, is_nan) {
            (_, _, true) => Err(StatsError::BadParams),
            (_, _, false) if shape_a <= 0.0 || shape_b <= 0.0 => Err(StatsError::BadParams),
            (_, _, false) => {
                Ok(Beta {
                    shape_a: shape_a,
                    shape_b: shape_b,
                })
            }
        }
    }

    /// Returns the shapeA (α) of the beta distribution
    ///
    /// # Examples
    ///
    /// ```
    /// use statrs::distribution::Beta;
    ///
    /// let n = Beta::new(2.0, 2.0).unwrap();
    /// assert_eq!(n.shape_a(), 2.0);
    /// ```
    pub fn shape_a(&self) -> f64 {
        self.shape_a
    }

    /// Returns the shapeB (β) of the beta distributionβ
    ///
    /// # Examples
    ///
    /// ```
    /// use statrs::distribution::Beta;
    ///
    /// let n = Beta::new(2.0, 2.0).unwrap();
    /// assert_eq!(n.shape_b(), 2.0);
    /// ```
    pub fn shape_b(&self) -> f64 {
        self.shape_b
    }
}

impl Sample<f64> for Beta {
    /// Generate a random sample from a beta distribution
    /// using `r` as the source of randomness.
    /// Refer [here](#method.sample-1) for implementation details.
    fn sample<R: Rng>(&mut self, r: &mut R) -> f64 {
        super::Distribution::sample(self, r)
    }
}

impl IndependentSample<f64> for Beta {
    /// Generate a random independent sample from a beta distribution
    /// using `r` as the source of randomness.
    /// Refer [here](#method.sample-1) for implementation details.
    fn ind_sample<R: Rng>(&self, r: &mut R) -> f64 {
        super::Distribution::sample(self, r)
    }
}

impl Distribution<f64> for Beta {
    /// Generate a random sample from a beta distribution using
    /// `r` as the source of randomness. Generated by sampling
    /// two gamma distributions and normalizing.
    ///
    /// # Examples
    ///
    /// ```
    /// # extern crate rand;
    /// # extern crate statrs;
    /// use rand::StdRng;
    /// use statrs::distribution::{Beta, Distribution};
    ///
    /// # fn main() {
    /// let mut r = rand::StdRng::new().unwrap();
    /// let n = Beta::new(2.0, 2.0).unwrap();
    /// print!("{}", n.sample::<StdRng>(&mut r));
    /// # }
    /// ```
    fn sample<R: Rng>(&self, r: &mut R) -> f64 {
        let x = super::gamma::sample_unchecked(r, self.shape_a, 1.0);
        let y = super::gamma::sample_unchecked(r, self.shape_b, 1.0);
        x / (x + y)
    }
}

impl Univariate<f64, f64> for Beta {
    /// Calculates the cumulative distribution function for the beta distribution
    /// at `x`
    ///
    /// # Panics
    ///
    /// If `x < 0.0` or `x > 1.0`
    ///
    /// # Formula
    ///
    /// ```ignore
    /// I_x(α, β)
    /// ```
    ///
    /// where `α` is shapeA, `β` is shapeB, and `I_x` is the regularized
    /// lower incomplete beta function
    fn cdf(&self, x: f64) -> f64 {
        assert!(x >= 0.0 && x <= 1.0,
                format!("{}", StatsError::ArgIntervalIncl("x", 0.0, 1.0)));
        if x == 1.0 {
            1.0
        } else if self.shape_a == f64::INFINITY && self.shape_b == f64::INFINITY {
            if x < 0.5 { 0.0 } else { 1.0 }
        } else if self.shape_a == f64::INFINITY {
            if x < 1.0 { 0.0 } else { 1.0 }
        } else if self.shape_b == f64::INFINITY {
            0.0
        } else if self.shape_a == 1.0 && self.shape_b == 1.0 {
            x
        } else {
            beta::beta_reg(self.shape_a, self.shape_b, x)
        }
    }

    /// Returns the minimum value in the domain of the
    /// beta distribution representable by a double precision
    /// float
    ///
    /// # Formula
    ///
    /// ```ignore
    /// 0
    /// ```
    fn min(&self) -> f64 {
        0.0
    }

    /// Returns the maximum value in the domain of the
    /// beta distribution representable by a double precision
    /// float
    ///
    /// # Formula
    ///
    /// ```ignore
    /// 1
    /// ```
    fn max(&self) -> f64 {
        1.0
    }
}

impl Mean<f64, f64> for Beta {
    /// Returns the mean of the beta distribution
    ///
    /// # Formula
    ///
    /// ```ignore
    /// α / (α + β)
    /// ```
    ///
    /// where `α` is shapeA and `β` is shapeB
    fn mean(&self) -> f64 {
        if self.shape_a == f64::INFINITY && self.shape_b == f64::INFINITY {
            0.5
        } else if self.shape_a == f64::INFINITY {
            1.0
        } else if self.shape_b == f64::INFINITY {
            0.0
        } else {
            self.shape_a / (self.shape_a + self.shape_b)
        }
    }
}

impl Variance<f64, f64> for Beta {
    /// Returns the variance of the beta distribution
    ///
    /// # Formula
    ///
    /// ```ignore
    /// (α * β) / ((α + β)^2 * (α + β + 1))
    /// ```
    ///
    /// where `α` is shapeA and `β` is shapeB
    fn variance(&self) -> f64 {
        self.shape_a * self.shape_b /
        ((self.shape_a + self.shape_b) * (self.shape_a + self.shape_b) *
         (self.shape_a + self.shape_b + 1.0))
    }

    /// Returns the standard deviation of the beta distribution
    ///
    /// # Formula
    ///
    /// ```ignore
    /// sqrt((α * β) / ((α + β)^2 * (α + β + 1)))
    /// ```
    ///
    /// where `α` is shapeA and `β` is shapeB
    fn std_dev(&self) -> f64 {
        self.variance().sqrt()
    }
}

impl Entropy<f64> for Beta {
    /// Returns the entropy of the beta distribution
    ///
    /// # Formula
    ///
    /// ```ignore
    /// ln(B(α, β)) - (α - 1)ψ(α) - (β - 1)ψ(β) + (α + β - 2)ψ(α + β)
    /// ```
    ///
    /// where `α` is shapeA, `β` is shapeB and `ψ` is the digamma function
    fn entropy(&self) -> f64 {
        if self.shape_a == f64::INFINITY || self.shape_b == f64::INFINITY {
            0.0
        } else {
            beta::ln_beta(self.shape_a, self.shape_b) -
            (self.shape_a - 1.0) * gamma::digamma(self.shape_a) -
            (self.shape_b - 1.0) * gamma::digamma(self.shape_b) +
            (self.shape_a + self.shape_b - 2.0) * gamma::digamma(self.shape_a + self.shape_b)
        }
    }
}
<|endoftext|>


<commit_before>#![cfg_attr(feature = "nightly-testing", feature(plugin))]
#![cfg_attr(feature = "nightly-testing", plugin(clippy))]

#[macro_use]
extern crate clap;
extern crate rayon;
extern crate rustfmt;
extern crate serde;
#[macro_use]
extern crate serde_derive;
extern crate serde_json;
extern crate toml;
extern crate inflector;
#[macro_use]
extern crate lazy_static;
extern crate regex;
extern crate hyper;

mod cargo;
mod codegen;
mod config;
mod service;

use std::fs::{self, OpenOptions};
use std::io::{Write, BufWriter};
use std::path::Path;

use rayon::prelude::*;

use clap::{Arg, App};

use codegen::botocore::{ServiceDefinition};
pub use service::Service;
use config::ServiceConfig;

fn main() {
    let matches = App::new("Rusoto Service Crate Generator")
        .version(crate_version!())
        .author(crate_authors!())
        .about(crate_description!())
        .arg(Arg::with_name("services_config")
            .long("config")
            .short("c")
            .takes_value(true))
        .arg(Arg::with_name("out_dir")
            .long("outdir")
            .short("o")
            .takes_value(true)
            .required(true))
        .get_matches();

    let out_dir = Path::new(matches.value_of("out_dir").unwrap());

    if !out_dir.exists() {
        fs::create_dir(out_dir).expect("Unable to create output directory");
    }

    let services_config_path = matches.value_of("services_config").unwrap();
    let service_configs = ServiceConfig::load_all(services_config_path).expect("Unable to read services configuration file.");

    service_configs.par_iter().for_each(|&(ref name, ref service_config)| {
        let service = {
            let service_definition = ServiceDefinition::load(name, &service_config.protocol_version)
                .expect(&format!("Failed to load service {}. Make sure the botocore submodule has been initialized!", name));
            Service::new(service_config.clone(), service_definition)
        };

        let crate_dir = out_dir.join(&name);
        let crate_name = format!("rusoto_{}", &name.replace('-', "_"));

        println!("Generating crate for {} @ {}...", service.full_name(), service.api_version());

        if !crate_dir.exists() {
            fs::create_dir(&crate_dir).expect(&format!("Unable to create directory at {}", crate_dir.display()));
        }

        let service_dependencies = service.get_dependencies();

        let extern_crates = service_dependencies.iter().map(|(k, _)| {
            if k == "xml-rs" {
                return "extern crate xml;".into();
            }
            let safe_name = k.replace("-", "_");
            let use_macro = k == "serde_derive" || k == "log" || k == "lazy_static";
            if use_macro {
                return format!("#[macro_use]\nextern crate {};", safe_name);
            }
            format!("extern crate {};", safe_name)
        }).collect::<Vec<String>>().join("\n");

        let mut cargo_manifest = OpenOptions::new()
            .write(true)
            .truncate(true)
            .create(true)
            .open(crate_dir.join("Cargo.toml"))
            .expect("Unable to write Cargo.toml");

        // crates.io has a max of 20 chars for keywords, do an ugly truncate if needed:
        let mut name_for_keyword = name.clone().to_string();
        if name_for_keyword.len() >= 20 {
            name_for_keyword = name_for_keyword[..20].to_string();
        }

        let manifest = cargo::Manifest {
            package: cargo::Metadata {
                authors: Some(vec![
                    "Anthony DiMarco <ocramida@gmail.com>".into(),
                    "Jimmy Cuadra <jimmy@jimmycuadra.com>".into(),
                    "Matthew Mayer <matthewkmayer@gmail.com>".into(),
                    "Nikita Pekin <contact@nikitapek.in>".into()
                ]),
                description: Some(format!("AWS SDK for Rust - {} @ {}", service.full_name(), service.api_version())),
                documentation: Some("https://rusoto.github.io/rusoto/rusoto_core/index.html".into()),
                keywords: Some(vec!["AWS".into(), "Amazon".into(), name_for_keyword]),
                license: Some("MIT".into()),
                name: crate_name.clone(),
                readme: Some("README.md".into()),
                repository: Some("https://github.com/rusoto/rusoto".into()),
                version: service_config.version.clone(),
                homepage: Some("https://www.rusoto.org/".into()),
                ..cargo::Metadata::default()
            },
            dependencies: service_dependencies,
            dev_dependencies: vec![
                ("rusoto_mock".to_owned(), cargo::Dependency::Extended {
                    path: Some("../../../mock".into()),
                    version: Some("0.24.0".into()),
                    optional: None,
                    default_features: None,
                    features: None
                })
            ].into_iter().collect(),
            ..cargo::Manifest::default()
        };

        cargo_manifest.write_all(toml::to_string(&manifest).unwrap().as_bytes()).unwrap();

        let mut readme_file = OpenOptions::new()
            .write(true)
            .truncate(true)
            .create(true)
            .open(crate_dir.join("README.md"))
            .expect("Unable to write README.md");

        let readme = format!(r#"
# Rusoto {short_name}
Rust SDK for {aws_name}

You may be looking for:

* [An overview of Rusoto][rusoto-overview]
* [AWS services supported by Rusoto][supported-aws-services]
* [API documentation][api-documentation]
* [Getting help with Rusoto][rusoto-help]

## Requirements

Rust stable or beta are required to use Rusoto. Nightly is tested, but not guaranteed to be supported. Older
versions _may_ be supported. The currently supported Rust versions can be found in the Rusoto project
[`travis.yml`](https://github.com/rusoto/rusoto/blob/master/.travis.yml).

On Linux, OpenSSL is required.

## Installation

To use `{crate_name}` in your application, add it as a dependency in your `Cargo.toml`:

```toml
[dependencies]
{crate_name} = "{version}"
```

## Contributing

See [CONTRIBUTING][contributing].

## License

Rusoto is distributed under the terms of the MIT license.

See [LICENSE][license] for details.

[api-documentation]: https://rusoto.github.io/rusoto/rusoto/ "API documentation"
[license]: https://github.com/rusoto/rusoto/blob/master/LICENSE "MIT License"
[contributing]: https://github.com/rusoto/rusoto/blob/master/CONTRIBUTING.md "Contributing Guide"
[rusoto-help]: https://www.rusoto.org/help.html "Getting help with Rusoto"
[rusoto-overview]: https://www.rusoto.org/ "Rusoto overview"
[supported-aws-services]: https://www.rusoto.org/supported-aws-services.html "List of AWS services supported by Rusoto"
        "#,
        short_name = service.service_type_name(),
        aws_name = service.full_name(),
        crate_name = crate_name,
        version = service_config.version
        );

        readme_file.write_all(readme.as_bytes()).unwrap();

        {
            let src_dir = crate_dir.join("src");

            if !src_dir.exists() {
                fs::create_dir(&src_dir).expect(&format!("Unable to create directory at {}", src_dir.display()));
            }

            let lib_file_path = src_dir.join("lib.rs");

            let mut lib_file = OpenOptions::new()
                .write(true)
                .truncate(true)
                .create(true)
                .open(&lib_file_path)
                .expect("Unable to write lib.rs");

            let lib_file_contents = format!(r#"
//! {service_full_name}
//!
//! If you're using the service, you're probably looking for [{client_name}](struct.{client_name}.html) and [{trait_name}](trait.{trait_name}.html).

{extern_crates}

mod generated;
mod custom;

pub use generated::*;
pub use custom::*;
            "#,
            service_full_name = service.full_name(),
            client_name = service.client_type_name(),
            trait_name = service.service_type_name(),
            extern_crates = extern_crates
            );

            lib_file.write_all(lib_file_contents.as_bytes()).unwrap();

            let gen_file_path = src_dir.join("generated.rs");

            let gen_file = OpenOptions::new()
                .create(true)
                .write(true)
                .truncate(true)
                .open(&gen_file_path)
                .expect("Unable to write generated.rs");

            let mut gen_writer = BufWriter::new(gen_file);

            codegen::generator::generate_source(&service, &mut gen_writer).unwrap();

            let custom_dir_path = src_dir.join("custom");

            if !custom_dir_path.exists() {
                fs::create_dir(&custom_dir_path).expect(&format!("Unable to create directory at {}", custom_dir_path.display()));
            }

            let custom_mod_file_path = custom_dir_path.join("mod.rs");

            if !custom_mod_file_path.exists() {
                OpenOptions::new()
                    .write(true)
                    .create_new(true)
                    .open(&custom_mod_file_path)
                    .expect("Unable to write mod.rs");
            }
        }

        {
            let src_dir = crate_dir.join("src");
            let gen_file_path = src_dir.join("generated.rs");

            let _ = rustfmt::run(rustfmt::Input::File(gen_file_path), &rustfmt::config::Config {
                write_mode: rustfmt::config::WriteMode::Overwrite,
                error_on_line_overflow: false,
                ..rustfmt::config::Config::default()
            });
        }

        {
            let test_resources_dir = crate_dir.join("test_resources");

            if !test_resources_dir.exists() {
                fs::create_dir(&test_resources_dir).expect(&format!("Unable to create directory at {}", test_resources_dir.display()));
            }

            let generated_test_resources_dir = test_resources_dir.join("generated");

            if !generated_test_resources_dir.exists() {
                fs::create_dir(&generated_test_resources_dir).expect(&format!("Unable to create directory at {}", generated_test_resources_dir.display()));
            }

            let test_valid_resources = codegen::generator::tests::find_valid_responses_for_service(&service);
            if !test_valid_resources.is_empty() {
                let test_valid_resources_dir = generated_test_resources_dir.join("valid");

                if !test_valid_resources_dir.exists() {
                    fs::create_dir(&test_valid_resources_dir).expect(&format!("Unable to create directory at {}", generated_test_resources_dir.display()));
                }

                for resource in test_valid_resources {
                    fs::copy(resource.full_path, test_valid_resources_dir.join(&resource.file_name)).expect("Failed to copy test resource file");
                }
            }

            let test_error_resources = codegen::generator::tests::find_error_responses_for_service(&service);
            if !test_error_resources.is_empty() {
                let test_error_resources_dir = generated_test_resources_dir.join("error");

                if !test_error_resources_dir.exists() {
                    fs::create_dir(&test_error_resources_dir).expect(&format!("Unable to create directory at {}", generated_test_resources_dir.display()));
                }

                for resource in test_error_resources {
                    fs::copy(resource.full_path, test_error_resources_dir.join(&resource.file_name)).expect("Failed to copy test resource file");
                }
            }
        }
    });
}
<commit_msg>tweak service_crategen to generate Cargo.toml specifying rusoto_mock 0.25.0 for service crates<commit_after>#![cfg_attr(feature = "nightly-testing", feature(plugin))]
#![cfg_attr(feature = "nightly-testing", plugin(clippy))]

#[macro_use]
extern crate clap;
extern crate rayon;
extern crate rustfmt;
extern crate serde;
#[macro_use]
extern crate serde_derive;
extern crate serde_json;
extern crate toml;
extern crate inflector;
#[macro_use]
extern crate lazy_static;
extern crate regex;
extern crate hyper;

mod cargo;
mod codegen;
mod config;
mod service;

use std::fs::{self, OpenOptions};
use std::io::{Write, BufWriter};
use std::path::Path;

use rayon::prelude::*;

use clap::{Arg, App};

use codegen::botocore::{ServiceDefinition};
pub use service::Service;
use config::ServiceConfig;

fn main() {
    let matches = App::new("Rusoto Service Crate Generator")
        .version(crate_version!())
        .author(crate_authors!())
        .about(crate_description!())
        .arg(Arg::with_name("services_config")
            .long("config")
            .short("c")
            .takes_value(true))
        .arg(Arg::with_name("out_dir")
            .long("outdir")
            .short("o")
            .takes_value(true)
            .required(true))
        .get_matches();

    let out_dir = Path::new(matches.value_of("out_dir").unwrap());

    if !out_dir.exists() {
        fs::create_dir(out_dir).expect("Unable to create output directory");
    }

    let services_config_path = matches.value_of("services_config").unwrap();
    let service_configs = ServiceConfig::load_all(services_config_path).expect("Unable to read services configuration file.");

    service_configs.par_iter().for_each(|&(ref name, ref service_config)| {
        let service = {
            let service_definition = ServiceDefinition::load(name, &service_config.protocol_version)
                .expect(&format!("Failed to load service {}. Make sure the botocore submodule has been initialized!", name));
            Service::new(service_config.clone(), service_definition)
        };

        let crate_dir = out_dir.join(&name);
        let crate_name = format!("rusoto_{}", &name.replace('-', "_"));

        println!("Generating crate for {} @ {}...", service.full_name(), service.api_version());

        if !crate_dir.exists() {
            fs::create_dir(&crate_dir).expect(&format!("Unable to create directory at {}", crate_dir.display()));
        }

        let service_dependencies = service.get_dependencies();

        let extern_crates = service_dependencies.iter().map(|(k, _)| {
            if k == "xml-rs" {
                return "extern crate xml;".into();
            }
            let safe_name = k.replace("-", "_");
            let use_macro = k == "serde_derive" || k == "log" || k == "lazy_static";
            if use_macro {
                return format!("#[macro_use]\nextern crate {};", safe_name);
            }
            format!("extern crate {};", safe_name)
        }).collect::<Vec<String>>().join("\n");

        let mut cargo_manifest = OpenOptions::new()
            .write(true)
            .truncate(true)
            .create(true)
            .open(crate_dir.join("Cargo.toml"))
            .expect("Unable to write Cargo.toml");

        // crates.io has a max of 20 chars for keywords, do an ugly truncate if needed:
        let mut name_for_keyword = name.clone().to_string();
        if name_for_keyword.len() >= 20 {
            name_for_keyword = name_for_keyword[..20].to_string();
        }

        let manifest = cargo::Manifest {
            package: cargo::Metadata {
                authors: Some(vec![
                    "Anthony DiMarco <ocramida@gmail.com>".into(),
                    "Jimmy Cuadra <jimmy@jimmycuadra.com>".into(),
                    "Matthew Mayer <matthewkmayer@gmail.com>".into(),
                    "Nikita Pekin <contact@nikitapek.in>".into()
                ]),
                description: Some(format!("AWS SDK for Rust - {} @ {}", service.full_name(), service.api_version())),
                documentation: Some("https://rusoto.github.io/rusoto/rusoto_core/index.html".into()),
                keywords: Some(vec!["AWS".into(), "Amazon".into(), name_for_keyword]),
                license: Some("MIT".into()),
                name: crate_name.clone(),
                readme: Some("README.md".into()),
                repository: Some("https://github.com/rusoto/rusoto".into()),
                version: service_config.version.clone(),
                homepage: Some("https://www.rusoto.org/".into()),
                ..cargo::Metadata::default()
            },
            dependencies: service_dependencies,
            dev_dependencies: vec![
                ("rusoto_mock".to_owned(), cargo::Dependency::Extended {
                    path: Some("../../../mock".into()),
                    version: Some("0.25.0".into()),
                    optional: None,
                    default_features: None,
                    features: None
                })
            ].into_iter().collect(),
            ..cargo::Manifest::default()
        };

        cargo_manifest.write_all(toml::to_string(&manifest).unwrap().as_bytes()).unwrap();

        let mut readme_file = OpenOptions::new()
            .write(true)
            .truncate(true)
            .create(true)
            .open(crate_dir.join("README.md"))
            .expect("Unable to write README.md");

        let readme = format!(r#"
# Rusoto {short_name}
Rust SDK for {aws_name}

You may be looking for:

* [An overview of Rusoto][rusoto-overview]
* [AWS services supported by Rusoto][supported-aws-services]
* [API documentation][api-documentation]
* [Getting help with Rusoto][rusoto-help]

## Requirements

Rust stable or beta are required to use Rusoto. Nightly is tested, but not guaranteed to be supported. Older
versions _may_ be supported. The currently supported Rust versions can be found in the Rusoto project
[`travis.yml`](https://github.com/rusoto/rusoto/blob/master/.travis.yml).

On Linux, OpenSSL is required.

## Installation

To use `{crate_name}` in your application, add it as a dependency in your `Cargo.toml`:

```toml
[dependencies]
{crate_name} = "{version}"
```

## Contributing

See [CONTRIBUTING][contributing].

## License

Rusoto is distributed under the terms of the MIT license.

See [LICENSE][license] for details.

[api-documentation]: https://rusoto.github.io/rusoto/rusoto/ "API documentation"
[license]: https://github.com/rusoto/rusoto/blob/master/LICENSE "MIT License"
[contributing]: https://github.com/rusoto/rusoto/blob/master/CONTRIBUTING.md "Contributing Guide"
[rusoto-help]: https://www.rusoto.org/help.html "Getting help with Rusoto"
[rusoto-overview]: https://www.rusoto.org/ "Rusoto overview"
[supported-aws-services]: https://www.rusoto.org/supported-aws-services.html "List of AWS services supported by Rusoto"
        "#,
        short_name = service.service_type_name(),
        aws_name = service.full_name(),
        crate_name = crate_name,
        version = service_config.version
        );

        readme_file.write_all(readme.as_bytes()).unwrap();

        {
            let src_dir = crate_dir.join("src");

            if !src_dir.exists() {
                fs::create_dir(&src_dir).expect(&format!("Unable to create directory at {}", src_dir.display()));
            }

            let lib_file_path = src_dir.join("lib.rs");

            let mut lib_file = OpenOptions::new()
                .write(true)
                .truncate(true)
                .create(true)
                .open(&lib_file_path)
                .expect("Unable to write lib.rs");

            let lib_file_contents = format!(r#"
//! {service_full_name}
//!
//! If you're using the service, you're probably looking for [{client_name}](struct.{client_name}.html) and [{trait_name}](trait.{trait_name}.html).

{extern_crates}

mod generated;
mod custom;

pub use generated::*;
pub use custom::*;
            "#,
            service_full_name = service.full_name(),
            client_name = service.client_type_name(),
            trait_name = service.service_type_name(),
            extern_crates = extern_crates
            );

            lib_file.write_all(lib_file_contents.as_bytes()).unwrap();

            let gen_file_path = src_dir.join("generated.rs");

            let gen_file = OpenOptions::new()
                .create(true)
                .write(true)
                .truncate(true)
                .open(&gen_file_path)
                .expect("Unable to write generated.rs");

            let mut gen_writer = BufWriter::new(gen_file);

            codegen::generator::generate_source(&service, &mut gen_writer).unwrap();

            let custom_dir_path = src_dir.join("custom");

            if !custom_dir_path.exists() {
                fs::create_dir(&custom_dir_path).expect(&format!("Unable to create directory at {}", custom_dir_path.display()));
            }

            let custom_mod_file_path = custom_dir_path.join("mod.rs");

            if !custom_mod_file_path.exists() {
                OpenOptions::new()
                    .write(true)
                    .create_new(true)
                    .open(&custom_mod_file_path)
                    .expect("Unable to write mod.rs");
            }
        }

        {
            let src_dir = crate_dir.join("src");
            let gen_file_path = src_dir.join("generated.rs");

            let _ = rustfmt::run(rustfmt::Input::File(gen_file_path), &rustfmt::config::Config {
                write_mode: rustfmt::config::WriteMode::Overwrite,
                error_on_line_overflow: false,
                ..rustfmt::config::Config::default()
            });
        }

        {
            let test_resources_dir = crate_dir.join("test_resources");

            if !test_resources_dir.exists() {
                fs::create_dir(&test_resources_dir).expect(&format!("Unable to create directory at {}", test_resources_dir.display()));
            }

            let generated_test_resources_dir = test_resources_dir.join("generated");

            if !generated_test_resources_dir.exists() {
                fs::create_dir(&generated_test_resources_dir).expect(&format!("Unable to create directory at {}", generated_test_resources_dir.display()));
            }

            let test_valid_resources = codegen::generator::tests::find_valid_responses_for_service(&service);
            if !test_valid_resources.is_empty() {
                let test_valid_resources_dir = generated_test_resources_dir.join("valid");

                if !test_valid_resources_dir.exists() {
                    fs::create_dir(&test_valid_resources_dir).expect(&format!("Unable to create directory at {}", generated_test_resources_dir.display()));
                }

                for resource in test_valid_resources {
                    fs::copy(resource.full_path, test_valid_resources_dir.join(&resource.file_name)).expect("Failed to copy test resource file");
                }
            }

            let test_error_resources = codegen::generator::tests::find_error_responses_for_service(&service);
            if !test_error_resources.is_empty() {
                let test_error_resources_dir = generated_test_resources_dir.join("error");

                if !test_error_resources_dir.exists() {
                    fs::create_dir(&test_error_resources_dir).expect(&format!("Unable to create directory at {}", generated_test_resources_dir.display()));
                }

                for resource in test_error_resources {
                    fs::copy(resource.full_path, test_error_resources_dir.join(&resource.file_name)).expect("Failed to copy test resource file");
                }
            }
        }
    });
}
<|endoftext|>


<commit_before>//! Utilities for working with the WTF-8 encoding.
//!
//! WTF-8 is a hack that allows to represent potentially ill-formed UTF-16 byte
//! strings in a way somewhat compatible with UTF-8. Such bugs UTF-16 strings
//! can be found for example in Windows paths for legacy reasons.
//!
//! See the official [WTF-8][wtf8] specification for more information.
//!
//! [wtf8]: https://simonsapin.github.io/wtf-8

/// Converts the given potentially ill-formed UTF-16 into a WTF-8 byte sequence.
///
/// If the given byte sequence is a well-formed UTF-16 string, then the result
/// is guaranteed to be a valid UTF-8string.
pub fn from_ill_formed_utf16(units: impl Iterator<Item=u16>) -> Vec<u8> {
    let mut res = Vec::new();

    let mut iter = units.peekable();
    loop {
        let unit = match iter.next() {
            Some(unit) => unit,
            None => return res,
        };

        // Potentially ill-formed UTF-16 to code point conversion based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#decoding-ill-formed-utf-16
        let mut point = unit as u32;
        if is_lead_surrogate(&unit) {
            if let Some(next) = iter.next_if(is_trail_surrogate) {
                let lead = (unit as u32 - 0xD800) << 10;
                let trail = next as u32 - 0xDC00;
                point = 0x10000 + lead + trail;
            }
        }

        // Code point to WTF-8 encoding procedure based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#encoding-wtf-8
        match point {
            0x0000..=0x007F => {
                res.push(point as u8);
            }
            0x0080..=0x07FF => {
                res.push(0xC0 | (point >> 6) as u8);
                res.push(0x80 | (point & 0x3F) as u8);
            }
            0x0800..=0xFFFF => {
                res.push(0xE0 | (point >> 12) as u8);
                res.push(0x80 | ((point >> 6) as u8 & 0x3F));
                res.push(0x80 | (point & 0x3F) as u8);
            }
            0x10000..=0x10FFFF => {
                res.push(0xF0 | (point >> 18) as u8);
                res.push(0x80 | ((point >> 12) as u8 & 0x3F));
                res.push(0x80 | ((point >> 6) as u8 & 0x3F));
                res.push(0x80 | (point & 0x3F) as u8);
            }
            _ => panic!(), // Not possible by construction.
        }
    }
}

/// Converts the given WTF-8 byte sequence into potentially ill-formed UTF-16.
///
/// If the given byte sequence is a valid UTF-8 string, then the result is
/// guaranteed to be a well-formed UTF-16 string.
pub fn into_ill_formed_utf16(units: impl Iterator<Item=u8>) -> Vec<u16> {
    let mut res = Vec::new();

    let mut iter = units.peekable();
    loop {
        let byte1 = match iter.next() {
            Some(byte1) => byte1,
            None => return res,
        };

        // TODO: Add support for proper error handling.

        // WTF-8 to a code point decoding procedure based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#decoding-wtf-8
        let mut point = 0;
        match byte1 {
            0x00..=0x7F => {
                point += byte1 as u32;
            }
            0xC2..=0xDF => {
                let byte2 = iter.next().unwrap_or_else(|| todo!());
                point += ((byte1 & 0x1F) as u32) << 6;
                point += ((byte2 & 0x3F) as u32) << 0;
            }
            0xE0..=0xEF => {
                let byte2 = iter.next().unwrap_or_else(|| todo!());
                let byte3 = iter.next().unwrap_or_else(|| todo!());
                point += ((byte1 & 0x0F) as u32) << 12;
                point += ((byte2 & 0x3F) as u32) << 6;
                point += ((byte3 & 0x3F) as u32) << 0;
            }
            0xF0..=0xF4 => {
                let byte2 = iter.next().unwrap_or_else(|| todo!());
                let byte3 = iter.next().unwrap_or_else(|| todo!());
                let byte4 = iter.next().unwrap_or_else(|| todo!());
                point += ((byte1 & 0x07) as u32) << 18;
                point += ((byte2 & 0x3F) as u32) << 12;
                point += ((byte3 & 0x3F) as u32) << 6;
                point += ((byte4 & 0x3F) as u32) << 0;
            }
            _ => todo!(),
        }

        // Code point to potentially ill-formed UTF-16 coversion based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#encoding-ill-formed-utf-16
        if is_supplementary(&point) {
            res.push(((point - 0x10000) >> 10) as u16 + 0xD800);
            res.push(((point - 0x10000) as u16 & 0x3FF) + 0xDC00);
        } else {
            res.push(point as u16);
        }
    }
}

/// Determines whether the given UTF-16 code unit is a lead surrogate.
#[inline]
fn is_lead_surrogate(unit: &u16) -> bool {
    matches!(unit, 0xD800..=0xDBFF)
}

/// Determines whether the given UTF-16 code unit is a trail surrogate.
#[inline]
fn is_trail_surrogate(unit: &u16) -> bool {
    matches!(unit, 0xDC00..=0xDFFF)
}

/// Determines whether the given code point is in the supplementary.
#[inline]
fn is_supplementary(point: &u32) -> bool {
    matches!(point, 0x10000..=0x10FFFF)
}

#[cfg(test)]
mod tests {

    use super::*;

    use quickcheck_macros::quickcheck;

    #[test]
    fn from_empty() {
        assert_eq!(from_ill_formed_utf16(std::iter::empty()), b"");
    }

    #[test]
    fn from_small_bytes() {
        let bytes_utf16 = vec![0x04, 0x08, 0x16, 0x23, 0x42];
        let bytes_wtf8 = from_ill_formed_utf16(bytes_utf16.into_iter());
        assert_eq!(bytes_wtf8, b"\x04\x08\x16\x23\x42");
    }

    #[test]
    fn from_bmp_bytes() {
        let bytes_utf16 = vec![0xABCD, 0xBEEF, 0xFFFF];
        let bytes_wtf8 = from_ill_formed_utf16(bytes_utf16.into_iter());
        assert_eq!(bytes_wtf8, b"\xEA\xAF\x8D\xEB\xBB\xAF\xEF\xBF\xBF");
    }

    #[test]
    fn from_sup_bytes() {
        let bytes_utf16 = vec![0xD800, 0xDC00, 0xDBFF, 0xDFFF];
        let bytes_wtf8 = from_ill_formed_utf16(bytes_utf16.into_iter());
        assert_eq!(bytes_wtf8, b"\xF0\x90\x80\x80\xF4\x8F\xBF\xBF");
    }

    #[test]
    fn from_ascii_string() {
        let string = "foo bar baz";
        let string_wtf8 = from_ill_formed_utf16(string.encode_utf16());
        let string_utf8 = string.as_bytes();
        assert_eq!(string_wtf8, string_utf8);
    }

    #[test]
    fn from_unicode_string() {
        let string = "zażółć gęślą jaźń";
        let string_wtf8 = from_ill_formed_utf16(string.encode_utf16());
        let string_utf8 = string.as_bytes();
        assert_eq!(string_wtf8, string_utf8);
    }

    #[test]
    fn into_empty() {
        assert_eq!(into_ill_formed_utf16(std::iter::empty()), vec![]);
    }

    #[test]
    fn into_ascii_string() {
        let string = "foo bar baz";
        let string_wtf8 = into_ill_formed_utf16(string.bytes());
        let string_utf16 = string.encode_utf16().collect::<Vec<_>>();
        assert_eq!(string_wtf8, string_utf16);
    }

    #[test]
    fn into_unicode_string() {
        let string = "zażółć gęślą jaźń";
        let string_wtf8 = into_ill_formed_utf16(string.bytes());
        let string_utf16 = string.encode_utf16().collect::<Vec<_>>();
        assert_eq!(string_wtf8, string_utf16);
    }

    #[quickcheck]
    fn from_any_string(input: String) {
        let units = input.encode_utf16();

        let string_wtf8 = from_ill_formed_utf16(units);
        let string_utf8 = input.as_bytes();
        assert_eq!(string_wtf8, string_utf8);
    }

    #[quickcheck]
    fn into_from_any_string(input: String) {
        let units = input.bytes();

        let string_utf16 = into_ill_formed_utf16(units);
        let string_wtf8 = from_ill_formed_utf16(string_utf16.into_iter());
        assert_eq!(String::from_utf8(string_wtf8).unwrap(), input);
    }

    #[quickcheck]
    fn from_into_any_units(input: Vec<u16>) {
        let units = input.iter().map(|unit| *unit);

        let string_wtf8 = from_ill_formed_utf16(units);
        let string_utf16 = into_ill_formed_utf16(string_wtf8.into_iter());
        assert_eq!(string_utf16, input);
    }
}
<commit_msg>Use `unreachable!` macro for impossible code.<commit_after>//! Utilities for working with the WTF-8 encoding.
//!
//! WTF-8 is a hack that allows to represent potentially ill-formed UTF-16 byte
//! strings in a way somewhat compatible with UTF-8. Such bugs UTF-16 strings
//! can be found for example in Windows paths for legacy reasons.
//!
//! See the official [WTF-8][wtf8] specification for more information.
//!
//! [wtf8]: https://simonsapin.github.io/wtf-8

/// Converts the given potentially ill-formed UTF-16 into a WTF-8 byte sequence.
///
/// If the given byte sequence is a well-formed UTF-16 string, then the result
/// is guaranteed to be a valid UTF-8string.
pub fn from_ill_formed_utf16(units: impl Iterator<Item=u16>) -> Vec<u8> {
    let mut res = Vec::new();

    let mut iter = units.peekable();
    loop {
        let unit = match iter.next() {
            Some(unit) => unit,
            None => return res,
        };

        // Potentially ill-formed UTF-16 to code point conversion based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#decoding-ill-formed-utf-16
        let mut point = unit as u32;
        if is_lead_surrogate(&unit) {
            if let Some(next) = iter.next_if(is_trail_surrogate) {
                let lead = (unit as u32 - 0xD800) << 10;
                let trail = next as u32 - 0xDC00;
                point = 0x10000 + lead + trail;
            }
        }

        // Code point to WTF-8 encoding procedure based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#encoding-wtf-8
        match point {
            0x0000..=0x007F => {
                res.push(point as u8);
            }
            0x0080..=0x07FF => {
                res.push(0xC0 | (point >> 6) as u8);
                res.push(0x80 | (point & 0x3F) as u8);
            }
            0x0800..=0xFFFF => {
                res.push(0xE0 | (point >> 12) as u8);
                res.push(0x80 | ((point >> 6) as u8 & 0x3F));
                res.push(0x80 | (point & 0x3F) as u8);
            }
            0x10000..=0x10FFFF => {
                res.push(0xF0 | (point >> 18) as u8);
                res.push(0x80 | ((point >> 12) as u8 & 0x3F));
                res.push(0x80 | ((point >> 6) as u8 & 0x3F));
                res.push(0x80 | (point & 0x3F) as u8);
            }
            _ => unreachable!(), // Not possible by construction.
        }
    }
}

/// Converts the given WTF-8 byte sequence into potentially ill-formed UTF-16.
///
/// If the given byte sequence is a valid UTF-8 string, then the result is
/// guaranteed to be a well-formed UTF-16 string.
pub fn into_ill_formed_utf16(units: impl Iterator<Item=u8>) -> Vec<u16> {
    let mut res = Vec::new();

    let mut iter = units.peekable();
    loop {
        let byte1 = match iter.next() {
            Some(byte1) => byte1,
            None => return res,
        };

        // TODO: Add support for proper error handling.

        // WTF-8 to a code point decoding procedure based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#decoding-wtf-8
        let mut point = 0;
        match byte1 {
            0x00..=0x7F => {
                point += byte1 as u32;
            }
            0xC2..=0xDF => {
                let byte2 = iter.next().unwrap_or_else(|| todo!());
                point += ((byte1 & 0x1F) as u32) << 6;
                point += ((byte2 & 0x3F) as u32) << 0;
            }
            0xE0..=0xEF => {
                let byte2 = iter.next().unwrap_or_else(|| todo!());
                let byte3 = iter.next().unwrap_or_else(|| todo!());
                point += ((byte1 & 0x0F) as u32) << 12;
                point += ((byte2 & 0x3F) as u32) << 6;
                point += ((byte3 & 0x3F) as u32) << 0;
            }
            0xF0..=0xF4 => {
                let byte2 = iter.next().unwrap_or_else(|| todo!());
                let byte3 = iter.next().unwrap_or_else(|| todo!());
                let byte4 = iter.next().unwrap_or_else(|| todo!());
                point += ((byte1 & 0x07) as u32) << 18;
                point += ((byte2 & 0x3F) as u32) << 12;
                point += ((byte3 & 0x3F) as u32) << 6;
                point += ((byte4 & 0x3F) as u32) << 0;
            }
            _ => todo!(),
        }

        // Code point to potentially ill-formed UTF-16 coversion based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#encoding-ill-formed-utf-16
        if is_supplementary(&point) {
            res.push(((point - 0x10000) >> 10) as u16 + 0xD800);
            res.push(((point - 0x10000) as u16 & 0x3FF) + 0xDC00);
        } else {
            res.push(point as u16);
        }
    }
}

/// Determines whether the given UTF-16 code unit is a lead surrogate.
#[inline]
fn is_lead_surrogate(unit: &u16) -> bool {
    matches!(unit, 0xD800..=0xDBFF)
}

/// Determines whether the given UTF-16 code unit is a trail surrogate.
#[inline]
fn is_trail_surrogate(unit: &u16) -> bool {
    matches!(unit, 0xDC00..=0xDFFF)
}

/// Determines whether the given code point is in the supplementary.
#[inline]
fn is_supplementary(point: &u32) -> bool {
    matches!(point, 0x10000..=0x10FFFF)
}

#[cfg(test)]
mod tests {

    use super::*;

    use quickcheck_macros::quickcheck;

    #[test]
    fn from_empty() {
        assert_eq!(from_ill_formed_utf16(std::iter::empty()), b"");
    }

    #[test]
    fn from_small_bytes() {
        let bytes_utf16 = vec![0x04, 0x08, 0x16, 0x23, 0x42];
        let bytes_wtf8 = from_ill_formed_utf16(bytes_utf16.into_iter());
        assert_eq!(bytes_wtf8, b"\x04\x08\x16\x23\x42");
    }

    #[test]
    fn from_bmp_bytes() {
        let bytes_utf16 = vec![0xABCD, 0xBEEF, 0xFFFF];
        let bytes_wtf8 = from_ill_formed_utf16(bytes_utf16.into_iter());
        assert_eq!(bytes_wtf8, b"\xEA\xAF\x8D\xEB\xBB\xAF\xEF\xBF\xBF");
    }

    #[test]
    fn from_sup_bytes() {
        let bytes_utf16 = vec![0xD800, 0xDC00, 0xDBFF, 0xDFFF];
        let bytes_wtf8 = from_ill_formed_utf16(bytes_utf16.into_iter());
        assert_eq!(bytes_wtf8, b"\xF0\x90\x80\x80\xF4\x8F\xBF\xBF");
    }

    #[test]
    fn from_ascii_string() {
        let string = "foo bar baz";
        let string_wtf8 = from_ill_formed_utf16(string.encode_utf16());
        let string_utf8 = string.as_bytes();
        assert_eq!(string_wtf8, string_utf8);
    }

    #[test]
    fn from_unicode_string() {
        let string = "zażółć gęślą jaźń";
        let string_wtf8 = from_ill_formed_utf16(string.encode_utf16());
        let string_utf8 = string.as_bytes();
        assert_eq!(string_wtf8, string_utf8);
    }

    #[test]
    fn into_empty() {
        assert_eq!(into_ill_formed_utf16(std::iter::empty()), vec![]);
    }

    #[test]
    fn into_ascii_string() {
        let string = "foo bar baz";
        let string_wtf8 = into_ill_formed_utf16(string.bytes());
        let string_utf16 = string.encode_utf16().collect::<Vec<_>>();
        assert_eq!(string_wtf8, string_utf16);
    }

    #[test]
    fn into_unicode_string() {
        let string = "zażółć gęślą jaźń";
        let string_wtf8 = into_ill_formed_utf16(string.bytes());
        let string_utf16 = string.encode_utf16().collect::<Vec<_>>();
        assert_eq!(string_wtf8, string_utf16);
    }

    #[quickcheck]
    fn from_any_string(input: String) {
        let units = input.encode_utf16();

        let string_wtf8 = from_ill_formed_utf16(units);
        let string_utf8 = input.as_bytes();
        assert_eq!(string_wtf8, string_utf8);
    }

    #[quickcheck]
    fn into_from_any_string(input: String) {
        let units = input.bytes();

        let string_utf16 = into_ill_formed_utf16(units);
        let string_wtf8 = from_ill_formed_utf16(string_utf16.into_iter());
        assert_eq!(String::from_utf8(string_wtf8).unwrap(), input);
    }

    #[quickcheck]
    fn from_into_any_units(input: Vec<u16>) {
        let units = input.iter().map(|unit| *unit);

        let string_wtf8 = from_ill_formed_utf16(units);
        let string_utf16 = into_ill_formed_utf16(string_wtf8.into_iter());
        assert_eq!(string_utf16, input);
    }
}
<|endoftext|>


<commit_before>#![doc(test(attr(deny(warnings))))]
#![warn(missing_docs)]
#![cfg_attr(docsrs, feature(doc_cfg))]
#![cfg_attr(not(feature = "futures-v0_3"), allow(dead_code, unused_imports))]

//! A crate for integrating signal handling with the Tokio runtime.
//!
//! This provides the [`Signals`][Signals] struct which acts as a
//! [`Stream`][futures_core_0_3::stream::Stream] of signals.
//!
//! Note that the `futures-v0_3` feature of this crate must be
//! enabled for `Signals` to implement the `Stream` trait.
//!
//! # Example
//!
//! ```rust
//! use std::io::Error;
//!
//! use signal_hook::consts::signal::*;
//! use signal_hook_tokio::Signals;
//!
//! use futures::stream::StreamExt;
//!
//! async fn handle_signals(signals: Signals) {
//!     let mut signals = signals.fuse();
//!     while let Some(signal) = signals.next().await {
//!         match signal {
//!             SIGHUP => {
//!                 // Reload configuration
//!                 // Reopen the log file
//!             }
//!             SIGTERM | SIGINT | SIGQUIT => {
//!                 // Shutdown the system;
//!             },
//!             _ => unreachable!(),
//!         }
//!     }
//! }
//!
//! #[tokio::main]
//! async fn main() -> Result<(), Error> {
//!     let signals = Signals::new(&[
//!         SIGHUP,
//!         SIGTERM,
//!         SIGINT,
//!         SIGQUIT,
//!     ])?;
//!     let handle = signals.handle();
//!
//!     let signals_task = tokio::spawn(handle_signals(signals));
//!
//!     // Execute your main program logic
//!
//!     // Terminate the signal stream.
//!     handle.close();
//!     signals_task.await?;
//!
//!     Ok(())
//! }
//! ```

macro_rules! implement_signals_with_pipe {
    ($pipe:ty) => {
        use std::borrow::Borrow;
        use std::io::Error;

        use libc::c_int;

        pub use signal_hook::iterator::backend::Handle;
        use signal_hook::iterator::backend::{OwningSignalIterator, SignalDelivery};
        use signal_hook::iterator::exfiltrator::{Exfiltrator, SignalOnly};

        /// An asynchronous [`Stream`] of arriving signals.
        ///
        /// The stream doesn't return the signals in the order they were recieved by
        /// the process and may merge signals received multiple times.
        pub struct SignalsInfo<E: Exfiltrator = SignalOnly>(OwningSignalIterator<$pipe, E>);

        impl<E: Exfiltrator> SignalsInfo<E> {
            /// Create a `SignalsInfo` instance.
            ///
            /// This registers all the signals listed. The same restrictions (panics, errors) apply
            /// as with [`Handle::add_signal`].
            pub fn new<I, S>(signals: I) -> Result<Self, Error>
            where
                I: IntoIterator<Item = S>,
                S: Borrow<c_int>,
                E: Default,
            {
                Self::with_exfiltrator(signals, E::default())
            }

            /// A constructor with explicit exfiltrator.
            pub fn with_exfiltrator<I, S>(signals: I, exfiltrator: E) -> Result<Self, Error>
            where
                I: IntoIterator<Item = S>,
                S: Borrow<c_int>,
            {
                let (read, write) = UnixStream::pair()?;
                let inner = SignalDelivery::with_pipe(read, write, exfiltrator, signals)?;
                Ok(Self(OwningSignalIterator::new(inner)))
            }

            /// Get a shareable [`Handle`] for this `Signals` instance.
            ///
            /// This can be used to add further signals or close the [`Signals`] instance
            /// which terminates the whole signal stream.
            pub fn handle(&self) -> Handle {
                self.0.handle()
            }
        }

        /// Simplified version of the signals stream.
        ///
        /// This one simply returns the signal numbers, while [`SignalsInfo`] can provide additional
        /// information.
        pub type Signals = SignalsInfo<SignalOnly>;
    };
}

use tokio::net::UnixStream;
implement_signals_with_pipe!(UnixStream);

use std::pin::Pin;
use std::task::{Context, Poll};

#[cfg(feature = "futures-v0_3")]
use futures_core_0_3::Stream;

use signal_hook::iterator::backend::PollResult;
use tokio::io::{AsyncRead, ReadBuf};

impl Signals {
    fn has_signals(read: &mut UnixStream, ctx: &mut Context<'_>) -> Result<bool, Error> {
        let mut buf = [0u8];
        let mut read_buf = ReadBuf::new(&mut buf);
        match Pin::new(read).poll_read(ctx, &mut read_buf) {
            Poll::Pending => Ok(false),
            Poll::Ready(Ok(())) => Ok(true),
            Poll::Ready(Err(error)) => Err(error),
        }
    }
}

#[cfg_attr(docsrs, doc(cfg(feature = "futures-v0_3")))]
#[cfg(feature = "futures-v0_3")]
impl Stream for Signals {
    type Item = c_int;

    fn poll_next(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        match self.0.poll_signal(&mut |read| Self::has_signals(read, ctx)) {
            PollResult::Signal(sig) => Poll::Ready(Some(sig)),
            PollResult::Closed => Poll::Ready(None),
            PollResult::Pending => Poll::Pending,
            PollResult::Err(error) => panic!("Unexpected error: {}", error),
        }
    }
}
<commit_msg>Make tests pass again<commit_after>#![doc(test(attr(deny(warnings))))]
#![warn(missing_docs)]
#![cfg_attr(docsrs, feature(doc_cfg))]
#![cfg_attr(not(feature = "futures-v0_3"), allow(dead_code, unused_imports))]

//! A crate for integrating signal handling with the Tokio runtime.
//!
//! This provides the [`Signals`][Signals] struct which acts as a
//! [`Stream`] of signals.
//!
//! Note that the `futures-v0_3` feature of this crate must be
//! enabled for `Signals` to implement the `Stream` trait.
//!
//! # Example
//!
//! ```rust
//! # #[cfg(feature = "futures-v0_3")]
//! # mod test {
//! use std::io::Error;
//!
//! use signal_hook::consts::signal::*;
//! use signal_hook_tokio::Signals;
//!
//! use futures::stream::StreamExt;
//!
//! async fn handle_signals(signals: Signals) {
//!     let mut signals = signals.fuse();
//!     while let Some(signal) = signals.next().await {
//!         match signal {
//!             SIGHUP => {
//!                 // Reload configuration
//!                 // Reopen the log file
//!             }
//!             SIGTERM | SIGINT | SIGQUIT => {
//!                 // Shutdown the system;
//!             },
//!             _ => unreachable!(),
//!         }
//!     }
//! }
//!
//! #[tokio::main]
//! # pub
//! async fn main() -> Result<(), Error> {
//!     let signals = Signals::new(&[
//!         SIGHUP,
//!         SIGTERM,
//!         SIGINT,
//!         SIGQUIT,
//!     ])?;
//!     let handle = signals.handle();
//!
//!     let signals_task = tokio::spawn(handle_signals(signals));
//!
//!     // Execute your main program logic
//!
//!     // Terminate the signal stream.
//!     handle.close();
//!     signals_task.await?;
//!
//!     Ok(())
//! }
//! }
//! # fn main() -> Result<(), std::io::Error> {
//! #    #[cfg(feature = "futures-v0_3")]
//! #    test::main()?;
//! #    Ok(())
//! # }
//! ```

macro_rules! implement_signals_with_pipe {
    ($pipe:ty) => {
        use std::borrow::Borrow;
        use std::io::Error;

        use libc::c_int;

        pub use signal_hook::iterator::backend::Handle;
        use signal_hook::iterator::backend::{OwningSignalIterator, SignalDelivery};
        use signal_hook::iterator::exfiltrator::{Exfiltrator, SignalOnly};

        /// An asynchronous [`Stream`] of arriving signals.
        ///
        /// The stream doesn't return the signals in the order they were recieved by
        /// the process and may merge signals received multiple times.
        pub struct SignalsInfo<E: Exfiltrator = SignalOnly>(OwningSignalIterator<$pipe, E>);

        impl<E: Exfiltrator> SignalsInfo<E> {
            /// Create a `SignalsInfo` instance.
            ///
            /// This registers all the signals listed. The same restrictions (panics, errors) apply
            /// as with [`Handle::add_signal`].
            pub fn new<I, S>(signals: I) -> Result<Self, Error>
            where
                I: IntoIterator<Item = S>,
                S: Borrow<c_int>,
                E: Default,
            {
                Self::with_exfiltrator(signals, E::default())
            }

            /// A constructor with explicit exfiltrator.
            pub fn with_exfiltrator<I, S>(signals: I, exfiltrator: E) -> Result<Self, Error>
            where
                I: IntoIterator<Item = S>,
                S: Borrow<c_int>,
            {
                let (read, write) = UnixStream::pair()?;
                let inner = SignalDelivery::with_pipe(read, write, exfiltrator, signals)?;
                Ok(Self(OwningSignalIterator::new(inner)))
            }

            /// Get a shareable [`Handle`] for this `Signals` instance.
            ///
            /// This can be used to add further signals or close the [`Signals`] instance
            /// which terminates the whole signal stream.
            pub fn handle(&self) -> Handle {
                self.0.handle()
            }
        }

        /// Simplified version of the signals stream.
        ///
        /// This one simply returns the signal numbers, while [`SignalsInfo`] can provide additional
        /// information.
        pub type Signals = SignalsInfo<SignalOnly>;
    };
}

use tokio::net::UnixStream;
implement_signals_with_pipe!(UnixStream);

use std::pin::Pin;
use std::task::{Context, Poll};

#[cfg(feature = "futures-v0_3")]
use futures_core_0_3::Stream;

use signal_hook::iterator::backend::PollResult;
use tokio::io::{AsyncRead, ReadBuf};

impl Signals {
    fn has_signals(read: &mut UnixStream, ctx: &mut Context<'_>) -> Result<bool, Error> {
        let mut buf = [0u8];
        let mut read_buf = ReadBuf::new(&mut buf);
        match Pin::new(read).poll_read(ctx, &mut read_buf) {
            Poll::Pending => Ok(false),
            Poll::Ready(Ok(())) => Ok(true),
            Poll::Ready(Err(error)) => Err(error),
        }
    }
}

#[cfg_attr(docsrs, doc(cfg(feature = "futures-v0_3")))]
#[cfg(feature = "futures-v0_3")]
impl Stream for Signals {
    type Item = c_int;

    fn poll_next(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        match self.0.poll_signal(&mut |read| Self::has_signals(read, ctx)) {
            PollResult::Signal(sig) => Poll::Ready(Some(sig)),
            PollResult::Closed => Poll::Ready(None),
            PollResult::Pending => Poll::Pending,
            PollResult::Err(error) => panic!("Unexpected error: {}", error),
        }
    }
}
<|endoftext|>