<commit_before>use std::{
    any::{Any, TypeId},
    collections::{BTreeMap, HashMap},
    marker::PhantomData,
    sync::{mpsc, Arc, Mutex},
};

use crate::shell::WindowRequest;

use dces::entity::Entity;

/// Internal wrapper that stores a message inside a box.
#[derive(Debug)]
pub struct MessageBox {
    message: Box<dyn Any + Send>,
    message_type: TypeId,
    target: Entity,
}

impl MessageBox {
    /// Creates a new `MessageBox`.
    pub fn new<M: Any + Send>(message: M, target: Entity) -> Self {
        MessageBox {
            message: Box::new(message),
            target,
            message_type: TypeId::of::<M>(),
        }
    }

    /// Check if the given type is the type of the message.
    pub fn is_type<M: Any>(&self) -> bool {
        self.message_type == TypeId::of::<M>()
    }

    /// Returns the type of the event.
    pub fn message_type(&self) -> TypeId {
        self.message_type
    }

    /// Downcasts the box to an concrete message.
    pub fn downcast<M: Any>(self) -> Result<M, String> {
        if self.message_type == TypeId::of::<M>() {
            return Ok(*self.message.downcast::<M>().unwrap());
        }

        Err("Wrong message type".to_string())
    }

    /// Downcasts the box as reference of an concrete message.
    pub fn downcast_ref<M: Any>(&self) -> Result<&M, String> {
        if self.message_type == TypeId::of::<M>() {
            return Ok(&*self.message.downcast_ref::<M>().unwrap());
        }

        Err("Wrong message type".to_string())
    }
}

/// The `MessageAdapter` provides a thread save entry point to sent
/// and read messages inside widget entities. They are processed inside the
/// method `message` defined in each widgets `State` code.
///
/// # Example
///
#[doc="The following code example ..."]
/// ```rust
/// // State
/// fn say_hello(entity: Entity, message_adapter: MessageAdapter) {
///     message_adapter.send_message(String::from("Hello rustician"), entity);
///     message_adapter.send_message(String::from("Did you got my message?"), entity);
/// }
///
/// impl Mystate {}
///
/// // implementation fo the State
/// impl State for MyState {
///     fn message(&mut self, mut messages: MessageReader, _registry: &mut Registry, _ctx: &mut Context) {
///         for message in messages.read::<String>() {
///         println!("{}", message);
///     }
/// }
/// ```
#[doc="will print two lines to stdout:"]
#[doc="```text"]
#[doc="$ Hello rustician,"]
#[doc="$ Did you got my message?"]
#[derive(Clone, Debug)]
pub struct MessageAdapter {
    messages: Arc<Mutex<BTreeMap<Entity, HashMap<TypeId, Vec<MessageBox>>>>>,
    window_sender: mpsc::Sender<WindowRequest>,
}

impl MessageAdapter {
    /// Creates a new message adapter
    pub fn new(window_sender: mpsc::Sender<WindowRequest>) -> Self {
        MessageAdapter {
            messages: Arc::new(Mutex::new(BTreeMap::new())),
            window_sender,
        }
    }

    /// Send a new message to the message pipeline.
    pub fn send_message<M: Any + Send>(&self, message: M, target: Entity) {
        if !self.messages.lock().unwrap().contains_key(&target) {
            self.messages
                .lock()
                .expect("MessageAdapter::send_message: Cannot lock messages.")
                .insert(target, HashMap::new());
        }

        let type_id = TypeId::of::<M>();

        if !self
            .messages
            .lock()
            .expect("MessageAdapter::send_message: Cannot lock messages.")
            .get(&target)
            .unwrap()
            .contains_key(&type_id)
        {
            self.messages
                .lock()
                .expect("MessageAdapter::send_message: Cannot lock messages.")
                .get_mut(&target)
                .unwrap()
                .insert(type_id, vec![]);
        }

        self.messages
            .lock()
            .expect("MessageAdapter::send_message: Cannot lock messages.")
            .get_mut(&target)
            .unwrap()
            .get_mut(&type_id)
            .unwrap()
            .push(MessageBox::new(message, target));

        self.window_sender
            .send(WindowRequest::Redraw)
            .expect("MessageAdapter::send_message: Cannot send redraw request.");
    }

    /// Returns a list of entities that has messages.
    pub(crate) fn entities(&self) -> Vec<Entity> {
        self.messages
            .lock()
            .expect("MessageAdapter::entities: Cannot lock messages.")
            .keys()
            .cloned()
            .collect()
    }

    /// Removes all messages for the given target entity. This is used
    /// to remove messages for entities that does not have a `State`
    /// to read the messages.
    pub(crate) fn remove_message_for_entity(&self, target: Entity) {
        self.messages
            .lock()
            .expect("MessageAdapter::remove_message_for_entity: Cannot lock messages.")
            .remove(&target);
    }

    /// Returns the number of messages in the queue.
    pub fn len(&self) -> usize {
        self.messages
            .lock()
            .expect("MessageAdapter::len: Cannot lock messages.")
            .len()
    }

    /// Returns `true` if the event message contains no events.
    pub fn is_empty(&self) -> bool {
        self.messages
            .lock()
            .expect("MessageAdapter::is_empty: Cannot lock messages.")
            .is_empty()
    }

    /// Returns a message reader for the given entity. Moves all
    /// messages for the entity from the adapter to the reader.
    pub(crate) fn message_reader(&self, entity: Entity) -> MessageReader {
        let messages = if let Some(messages) = self
            .messages
            .lock()
            .expect("MessageAdapter::message_reader: Cannot lock messages.")
            .remove(&entity)
        {
            messages
        } else {
            HashMap::new()
        };

        MessageReader::new(messages, entity)
    }
}

/// The `MessageReader` is used to access the messages of a widget.
pub struct MessageReader {
    messages: HashMap<TypeId, Vec<MessageBox>>,
    target: Entity,
}

impl MessageReader {
    /// Creates a new message reader.
    pub fn new(messages: HashMap<TypeId, Vec<MessageBox>>, target: Entity) -> Self {
        MessageReader { messages, target }
    }

    /// Returns the target entity of the reader.
    pub fn entity(&self) -> Entity {
        self.target
    }

    /// Returns `true` if the reader contains message for the specified type.
    pub fn contains_type<M: Any>(&self) -> bool {
        self.messages.contains_key(&TypeId::of::<M>())
    }

    /// Returns `true` if the reader contains no messages.
    pub fn is_empty(&self) -> bool {
        self.messages.is_empty()
    }

    /// Returns a message iterator for the given message type.
    pub fn read<M: Any + Send>(&mut self) -> MessageIterator<M> {
        let messages = if let Some(messages) = self.messages.remove(&TypeId::of::<M>()) {
            messages
        } else {
            vec![]
        };

        MessageIterator::new(messages)
    }
}

/// Iterator of messages.
#[derive(Debug)]
pub struct MessageIterator<M>
where
    M: Any + Send,
{
    messages: Vec<MessageBox>,
    _phantom: PhantomData<M>,
}

impl<M> MessageIterator<M>
where
    M: Any + Send,
{
    pub(crate) fn new(messages: Vec<MessageBox>) -> Self {
        MessageIterator {
            messages,
            _phantom: PhantomData::default(),
        }
    }
}

impl<M> Iterator for MessageIterator<M>
where
    M: Any + Send,
{
    type Item = M;

    fn next(&mut self) -> Option<Self::Item> {
        if self.messages.is_empty() {
            return None;
        }

        // unwrap is ok because only messages of the same type should
        // be stored in the vec
        Some(self.messages.remove(0).downcast::<M>().unwrap())
    }
}
<commit_msg>orbtk_core: lint update<commit_after>use std::{
    any::{Any, TypeId},
    collections::{BTreeMap, HashMap},
    marker::PhantomData,
    sync::{mpsc, Arc, Mutex},
};

use crate::shell::WindowRequest;

use dces::entity::Entity;

/// Internal wrapper that stores a message inside a box.
#[derive(Debug)]
pub struct MessageBox {
    message: Box<dyn Any + Send>,
    message_type: TypeId,
    target: Entity,
}

impl MessageBox {
    /// Creates a new `MessageBox`.
    pub fn new<M: Any + Send>(message: M, target: Entity) -> Self {
        MessageBox {
            message: Box::new(message),
            target,
            message_type: TypeId::of::<M>(),
        }
    }

    /// Check if the given type is the type of the message.
    pub fn is_type<M: Any>(&self) -> bool {
        self.message_type == TypeId::of::<M>()
    }

    /// Returns the type of the event.
    pub fn message_type(&self) -> TypeId {
        self.message_type
    }

    /// Downcasts the box to an concrete message.
    pub fn downcast<M: Any>(self) -> Result<M, String> {
        if self.message_type == TypeId::of::<M>() {
            return Ok(*self.message.downcast::<M>().unwrap());
        }

        Err("Wrong message type".to_string())
    }

    /// Downcasts the box as reference of an concrete message.
    pub fn downcast_ref<M: Any>(&self) -> Result<&M, String> {
        if self.message_type == TypeId::of::<M>() {
            return Ok(&*self.message.downcast_ref::<M>().unwrap());
        }

        Err("Wrong message type".to_string())
    }
}

/// The `MessageAdapter` provides a thread save entry point to sent
/// and read messages inside widget entities. They are processed inside the
/// method `message` defined in each widgets `State` code.
///
/// # Example
///
#[doc = "The following code example ..."]
/// ```rust
/// // State
/// fn say_hello(entity: Entity, message_adapter: MessageAdapter) {
///     message_adapter.send_message(String::from("Hello rustician"), entity);
///     message_adapter.send_message(String::from("Did you got my message?"), entity);
/// }
///
/// impl Mystate {}
///
/// // implementation fo the State
/// impl State for MyState {
///     fn message(&mut self, mut messages: MessageReader, _registry: &mut Registry, _ctx: &mut Context) {
///         for message in messages.read::<String>() {
///         println!("{}", message);
///     }
/// }
/// ```
#[doc = "will print two lines to stdout:"]
#[doc = "```text"]
#[doc = "$ Hello rustician,"]
#[doc = "$ Did you got my message?"]
#[derive(Clone, Debug)]
pub struct MessageAdapter {
    messages: Arc<Mutex<BTreeMap<Entity, HashMap<TypeId, Vec<MessageBox>>>>>,
    window_sender: mpsc::Sender<WindowRequest>,
}

impl MessageAdapter {
    /// Creates a new message adapter
    pub fn new(window_sender: mpsc::Sender<WindowRequest>) -> Self {
        MessageAdapter {
            messages: Arc::new(Mutex::new(BTreeMap::new())),
            window_sender,
        }
    }

    /// Send a new message to the message pipeline.
    pub fn send_message<M: Any + Send>(&self, message: M, target: Entity) {
        if !self.messages.lock().unwrap().contains_key(&target) {
            self.messages
                .lock()
                .expect("MessageAdapter::send_message: Cannot lock messages.")
                .insert(target, HashMap::new());
        }

        let type_id = TypeId::of::<M>();

        if !self
            .messages
            .lock()
            .expect("MessageAdapter::send_message: Cannot lock messages.")
            .get(&target)
            .unwrap()
            .contains_key(&type_id)
        {
            self.messages
                .lock()
                .expect("MessageAdapter::send_message: Cannot lock messages.")
                .get_mut(&target)
                .unwrap()
                .insert(type_id, vec![]);
        }

        self.messages
            .lock()
            .expect("MessageAdapter::send_message: Cannot lock messages.")
            .get_mut(&target)
            .unwrap()
            .get_mut(&type_id)
            .unwrap()
            .push(MessageBox::new(message, target));

        self.window_sender
            .send(WindowRequest::Redraw)
            .expect("MessageAdapter::send_message: Cannot send redraw request.");
    }

    /// Returns a list of entities that has messages.
    pub(crate) fn entities(&self) -> Vec<Entity> {
        self.messages
            .lock()
            .expect("MessageAdapter::entities: Cannot lock messages.")
            .keys()
            .cloned()
            .collect()
    }

    /// Removes all messages for the given target entity. This is used
    /// to remove messages for entities that does not have a `State`
    /// to read the messages.
    pub(crate) fn remove_message_for_entity(&self, target: Entity) {
        self.messages
            .lock()
            .expect("MessageAdapter::remove_message_for_entity: Cannot lock messages.")
            .remove(&target);
    }

    /// Returns the number of messages in the queue.
    pub fn len(&self) -> usize {
        self.messages
            .lock()
            .expect("MessageAdapter::len: Cannot lock messages.")
            .len()
    }

    /// Returns `true` if the event message contains no events.
    pub fn is_empty(&self) -> bool {
        self.messages
            .lock()
            .expect("MessageAdapter::is_empty: Cannot lock messages.")
            .is_empty()
    }

    /// Returns a message reader for the given entity. Moves all
    /// messages for the entity from the adapter to the reader.
    pub(crate) fn message_reader(&self, entity: Entity) -> MessageReader {
        let messages = if let Some(messages) = self
            .messages
            .lock()
            .expect("MessageAdapter::message_reader: Cannot lock messages.")
            .remove(&entity)
        {
            messages
        } else {
            HashMap::new()
        };

        MessageReader::new(messages, entity)
    }
}

/// The `MessageReader` is used to access the messages of a widget.
pub struct MessageReader {
    messages: HashMap<TypeId, Vec<MessageBox>>,
    target: Entity,
}

impl MessageReader {
    /// Creates a new message reader.
    pub fn new(messages: HashMap<TypeId, Vec<MessageBox>>, target: Entity) -> Self {
        MessageReader { messages, target }
    }

    /// Returns the target entity of the reader.
    pub fn entity(&self) -> Entity {
        self.target
    }

    /// Returns `true` if the reader contains message for the specified type.
    pub fn contains_type<M: Any>(&self) -> bool {
        self.messages.contains_key(&TypeId::of::<M>())
    }

    /// Returns `true` if the reader contains no messages.
    pub fn is_empty(&self) -> bool {
        self.messages.is_empty()
    }

    /// Returns a message iterator for the given message type.
    pub fn read<M: Any + Send>(&mut self) -> MessageIterator<M> {
        let messages = if let Some(messages) = self.messages.remove(&TypeId::of::<M>()) {
            messages
        } else {
            vec![]
        };

        MessageIterator::new(messages)
    }
}

/// Iterator of messages.
#[derive(Debug)]
pub struct MessageIterator<M>
where
    M: Any + Send,
{
    messages: Vec<MessageBox>,
    _phantom: PhantomData<M>,
}

impl<M> MessageIterator<M>
where
    M: Any + Send,
{
    pub(crate) fn new(messages: Vec<MessageBox>) -> Self {
        MessageIterator {
            messages,
            _phantom: PhantomData::default(),
        }
    }
}

impl<M> Iterator for MessageIterator<M>
where
    M: Any + Send,
{
    type Item = M;

    fn next(&mut self) -> Option<Self::Item> {
        if self.messages.is_empty() {
            return None;
        }

        // unwrap is ok because only messages of the same type should
        // be stored in the vec
        Some(self.messages.remove(0).downcast::<M>().unwrap())
    }
}
<|endoftext|>


<commit_before>// Copyright 2013 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use extra::net::url::Url;

pub fn build_request(url: Url) -> ~str {

    let host = copy url.host;
    let mut path = if url.path.len() > 0 { copy url.path } else { ~"/" };

    if url.query.len() > 0 {
        let kvps = do url.query.map |pair| {
            match *pair {
                (ref key, ref value) => fmt!("%s=%s", *key, *value)
            }
        };
        path.push_str("?");
        path.push_str(kvps.connect("&"));
    }

    let request_header = fmt!("GET %s HTTP/1.0\u000D\u000AHost: %s\u000D\u000A\u000D\u000A",
                              path, host);

    return request_header;
}

#[test]
#[allow(non_implicitly_copyable_typarams)]
fn should_request_slash_when_path_is_empty() {
    use extra::net::url;
    let url = url::from_str("http://host").get();
    assert!(url.path.is_empty());
    let headers = build_request(url);
    assert!(headers.contains("GET / "));
}
<commit_msg>Escape query string values. Fixes Issue #16.<commit_after>// Copyright 2013 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use extra::net::url::Url;

pub fn build_request(url: Url) -> ~str {

    let host = copy url.host;
    let mut path = if url.path.len() > 0 { copy url.path } else { ~"/" };

    if url.query.len() > 0 {
        let kvps = do url.query.map |pair| {
            match *pair {
                (ref key, ref value) => {
                    let qs_value:~str = (*value).replace(" ", "%20");
                    fmt!("%s=%s", *key, qs_value)
                }
            }
        };
        path.push_str("?");
        path.push_str(kvps.connect("&"));
    }

    let request_header = fmt!("GET %s HTTP/1.0\u000D\u000AHost: %s\u000D\u000A\u000D\u000A",
                              path, host);

    return request_header;
}

#[test]
#[allow(non_implicitly_copyable_typarams)]
fn should_request_slash_when_path_is_empty() {
    use extra::net::url;
    let url = url::from_str("http://host").get();
    assert!(url.path.is_empty());
    let headers = build_request(url);
    assert!(headers.contains("GET / "));
}
<|endoftext|>


<commit_before>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The Gamma and derived distributions.

use core::num::Float;

use {Rng, Open01};
use super::normal::StandardNormal;
use super::{IndependentSample, Sample, Exp};

/// The Gamma distribution `Gamma(shape, scale)` distribution.
///
/// The density function of this distribution is
///
/// ```notrust
/// f(x) =  x^(k - 1) * exp(-x / θ) / (Γ(k) * θ^k)
/// ```
///
/// where `Γ` is the Gamma function, `k` is the shape and `θ` is the
/// scale and both `k` and `θ` are strictly positive.
///
/// The algorithm used is that described by Marsaglia & Tsang 2000[1],
/// falling back to directly sampling from an Exponential for `shape
/// == 1`, and using the boosting technique described in [1] for
/// `shape < 1`.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{IndependentSample, Gamma};
///
/// let gamma = Gamma::new(2.0, 5.0);
/// let v = gamma.ind_sample(&mut rand::task_rng());
/// println!("{} is from a Gamma(2, 5) distribution", v);
/// ```
///
/// [1]: George Marsaglia and Wai Wan Tsang. 2000. "A Simple Method
/// for Generating Gamma Variables" *ACM Trans. Math. Softw.* 26, 3
/// (September 2000),
/// 363-372. DOI:[10.1145/358407.358414](http://doi.acm.org/10.1145/358407.358414)
pub struct Gamma {
    repr: GammaRepr,
}

enum GammaRepr {
    Large(GammaLargeShape),
    One(Exp),
    Small(GammaSmallShape)
}

// These two helpers could be made public, but saving the
// match-on-Gamma-enum branch from using them directly (e.g. if one
// knows that the shape is always > 1) doesn't appear to be much
// faster.

/// Gamma distribution where the shape parameter is less than 1.
///
/// Note, samples from this require a compulsory floating-point `pow`
/// call, which makes it significantly slower than sampling from a
/// gamma distribution where the shape parameter is greater than or
/// equal to 1.
///
/// See `Gamma` for sampling from a Gamma distribution with general
/// shape parameters.
struct GammaSmallShape {
    inv_shape: f64,
    large_shape: GammaLargeShape
}

/// Gamma distribution where the shape parameter is larger than 1.
///
/// See `Gamma` for sampling from a Gamma distribution with general
/// shape parameters.
struct GammaLargeShape {
    shape: f64,
    scale: f64,
    c: f64,
    d: f64
}

impl Gamma {
    /// Construct an object representing the `Gamma(shape, scale)`
    /// distribution.
    ///
    /// Fails if `shape <= 0` or `scale <= 0`.
    pub fn new(shape: f64, scale: f64) -> Gamma {
        assert!(shape > 0.0, "Gamma::new called with shape <= 0");
        assert!(scale > 0.0, "Gamma::new called with scale <= 0");

        let repr = match shape {
            1.0        => One(Exp::new(1.0 / scale)),
            0.0 .. 1.0 => Small(GammaSmallShape::new_raw(shape, scale)),
            _          => Large(GammaLargeShape::new_raw(shape, scale))
        };
        Gamma { repr: repr }
    }
}

impl GammaSmallShape {
    fn new_raw(shape: f64, scale: f64) -> GammaSmallShape {
        GammaSmallShape {
            inv_shape: 1. / shape,
            large_shape: GammaLargeShape::new_raw(shape + 1.0, scale)
        }
    }
}

impl GammaLargeShape {
    fn new_raw(shape: f64, scale: f64) -> GammaLargeShape {
        let d = shape - 1. / 3.;
        GammaLargeShape {
            shape: shape,
            scale: scale,
            c: 1. / (9. * d).sqrt(),
            d: d
        }
    }
}

impl Sample<f64> for Gamma {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl Sample<f64> for GammaSmallShape {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl Sample<f64> for GammaLargeShape {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}

impl IndependentSample<f64> for Gamma {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        match self.repr {
            Small(ref g) => g.ind_sample(rng),
            One(ref g) => g.ind_sample(rng),
            Large(ref g) => g.ind_sample(rng),
        }
    }
}
impl IndependentSample<f64> for GammaSmallShape {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        let Open01(u) = rng.gen::<Open01<f64>>();

        self.large_shape.ind_sample(rng) * u.powf(self.inv_shape)
    }
}
impl IndependentSample<f64> for GammaLargeShape {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        loop {
            let StandardNormal(x) = rng.gen::<StandardNormal>();
            let v_cbrt = 1.0 + self.c * x;
            if v_cbrt <= 0.0 { // a^3 <= 0 iff a <= 0
                continue
            }

            let v = v_cbrt * v_cbrt * v_cbrt;
            let Open01(u) = rng.gen::<Open01<f64>>();

            let x_sqr = x * x;
            if u < 1.0 - 0.0331 * x_sqr * x_sqr ||
                u.ln() < 0.5 * x_sqr + self.d * (1.0 - v + v.ln()) {
                return self.d * v * self.scale
            }
        }
    }
}

/// The chi-squared distribution `χ²(k)`, where `k` is the degrees of
/// freedom.
///
/// For `k > 0` integral, this distribution is the sum of the squares
/// of `k` independent standard normal random variables. For other
/// `k`, this uses the equivalent characterisation `χ²(k) = Gamma(k/2,
/// 2)`.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{ChiSquared, IndependentSample};
///
/// let chi = ChiSquared::new(11.0);
/// let v = chi.ind_sample(&mut rand::task_rng());
/// println!("{} is from a χ²(11) distribution", v)
/// ```
pub struct ChiSquared {
    repr: ChiSquaredRepr,
}

enum ChiSquaredRepr {
    // k == 1, Gamma(alpha, ..) is particularly slow for alpha < 1,
    // e.g. when alpha = 1/2 as it would be for this case, so special-
    // casing and using the definition of N(0,1)^2 is faster.
    DoFExactlyOne,
    DoFAnythingElse(Gamma),
}

impl ChiSquared {
    /// Create a new chi-squared distribution with degrees-of-freedom
    /// `k`. Fails if `k < 0`.
    pub fn new(k: f64) -> ChiSquared {
        let repr = if k == 1.0 {
            DoFExactlyOne
        } else {
            assert!(k > 0.0, "ChiSquared::new called with `k` < 0");
            DoFAnythingElse(Gamma::new(0.5 * k, 2.0))
        };
        ChiSquared { repr: repr }
    }
}
impl Sample<f64> for ChiSquared {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl IndependentSample<f64> for ChiSquared {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        match self.repr {
            DoFExactlyOne => {
                // k == 1 => N(0,1)^2
                let StandardNormal(norm) = rng.gen::<StandardNormal>();
                norm * norm
            }
            DoFAnythingElse(ref g) => g.ind_sample(rng)
        }
    }
}

/// The Fisher F distribution `F(m, n)`.
///
/// This distribution is equivalent to the ratio of two normalised
/// chi-squared distributions, that is, `F(m,n) = (χ²(m)/m) /
/// (χ²(n)/n)`.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{FisherF, IndependentSample};
///
/// let f = FisherF::new(2.0, 32.0);
/// let v = f.ind_sample(&mut rand::task_rng());
/// println!("{} is from an F(2, 32) distribution", v)
/// ```
pub struct FisherF {
    numer: ChiSquared,
    denom: ChiSquared,
    // denom_dof / numer_dof so that this can just be a straight
    // multiplication, rather than a division.
    dof_ratio: f64,
}

impl FisherF {
    /// Create a new `FisherF` distribution, with the given
    /// parameter. Fails if either `m` or `n` are not positive.
    pub fn new(m: f64, n: f64) -> FisherF {
        assert!(m > 0.0, "FisherF::new called with `m < 0`");
        assert!(n > 0.0, "FisherF::new called with `n < 0`");

        FisherF {
            numer: ChiSquared::new(m),
            denom: ChiSquared::new(n),
            dof_ratio: n / m
        }
    }
}
impl Sample<f64> for FisherF {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl IndependentSample<f64> for FisherF {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        self.numer.ind_sample(rng) / self.denom.ind_sample(rng) * self.dof_ratio
    }
}

/// The Student t distribution, `t(nu)`, where `nu` is the degrees of
/// freedom.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{StudentT, IndependentSample};
///
/// let t = StudentT::new(11.0);
/// let v = t.ind_sample(&mut rand::task_rng());
/// println!("{} is from a t(11) distribution", v)
/// ```
pub struct StudentT {
    chi: ChiSquared,
    dof: f64
}

impl StudentT {
    /// Create a new Student t distribution with `n` degrees of
    /// freedom. Fails if `n <= 0`.
    pub fn new(n: f64) -> StudentT {
        assert!(n > 0.0, "StudentT::new called with `n <= 0`");
        StudentT {
            chi: ChiSquared::new(n),
            dof: n
        }
    }
}
impl Sample<f64> for StudentT {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl IndependentSample<f64> for StudentT {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        let StandardNormal(norm) = rng.gen::<StandardNormal>();
        norm * (self.dof / self.chi.ind_sample(rng)).sqrt()
    }
}

#[cfg(test)]
mod test {
    use std::prelude::*;

    use distributions::{Sample, IndependentSample};
    use super::{ChiSquared, StudentT, FisherF};

    #[test]
    fn test_chi_squared_one() {
        let mut chi = ChiSquared::new(1.0);
        let mut rng = ::test::rng();
        for _ in range(0, 1000) {
            chi.sample(&mut rng);
            chi.ind_sample(&mut rng);
        }
    }
    #[test]
    fn test_chi_squared_small() {
        let mut chi = ChiSquared::new(0.5);
        let mut rng = ::test::rng();
        for _ in range(0, 1000) {
            chi.sample(&mut rng);
            chi.ind_sample(&mut rng);
        }
    }
    #[test]
    fn test_chi_squared_large() {
        let mut chi = ChiSquared::new(30.0);
        let mut rng = ::test::rng();
        for _ in range(0, 1000) {
            chi.sample(&mut rng);
            chi.ind_sample(&mut rng);
        }
    }
    #[test]
    #[should_fail]
    fn test_chi_squared_invalid_dof() {
        ChiSquared::new(-1.0);
    }

    #[test]
    fn test_f() {
        let mut f = FisherF::new(2.0, 32.0);
        let mut rng = ::test::rng();
        for _ in range(0, 1000) {
            f.sample(&mut rng);
            f.ind_sample(&mut rng);
        }
    }

    #[test]
    fn test_t() {
        let mut t = StudentT::new(11.0);
        let mut rng = ::test::rng();
        for _ in range(0, 1000) {
            t.sample(&mut rng);
            t.ind_sample(&mut rng);
        }
    }
}

#[cfg(test)]
mod bench {
    extern crate test;
    use std::prelude::*;
    use self::test::Bencher;
    use std::mem::size_of;
    use distributions::IndependentSample;
    use super::Gamma;

    #[bench]
    fn bench_gamma_large_shape(b: &mut Bencher) {
        let gamma = Gamma::new(10., 1.0);
        let mut rng = ::test::weak_rng();

        b.iter(|| {
            for _ in range(0, ::RAND_BENCH_N) {
                gamma.ind_sample(&mut rng);
            }
        });
        b.bytes = size_of::<f64>() as u64 * ::RAND_BENCH_N;
    }

    #[bench]
    fn bench_gamma_small_shape(b: &mut Bencher) {
        let gamma = Gamma::new(0.1, 1.0);
        let mut rng = ::test::weak_rng();

        b.iter(|| {
            for _ in range(0, ::RAND_BENCH_N) {
                gamma.ind_sample(&mut rng);
            }
        });
        b.bytes = size_of::<f64>() as u64 * ::RAND_BENCH_N;
    }
}
<commit_msg>docs: Stop using `notrust`<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The Gamma and derived distributions.

use core::num::Float;

use {Rng, Open01};
use super::normal::StandardNormal;
use super::{IndependentSample, Sample, Exp};

/// The Gamma distribution `Gamma(shape, scale)` distribution.
///
/// The density function of this distribution is
///
/// ```text
/// f(x) =  x^(k - 1) * exp(-x / θ) / (Γ(k) * θ^k)
/// ```
///
/// where `Γ` is the Gamma function, `k` is the shape and `θ` is the
/// scale and both `k` and `θ` are strictly positive.
///
/// The algorithm used is that described by Marsaglia & Tsang 2000[1],
/// falling back to directly sampling from an Exponential for `shape
/// == 1`, and using the boosting technique described in [1] for
/// `shape < 1`.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{IndependentSample, Gamma};
///
/// let gamma = Gamma::new(2.0, 5.0);
/// let v = gamma.ind_sample(&mut rand::task_rng());
/// println!("{} is from a Gamma(2, 5) distribution", v);
/// ```
///
/// [1]: George Marsaglia and Wai Wan Tsang. 2000. "A Simple Method
/// for Generating Gamma Variables" *ACM Trans. Math. Softw.* 26, 3
/// (September 2000),
/// 363-372. DOI:[10.1145/358407.358414](http://doi.acm.org/10.1145/358407.358414)
pub struct Gamma {
    repr: GammaRepr,
}

enum GammaRepr {
    Large(GammaLargeShape),
    One(Exp),
    Small(GammaSmallShape)
}

// These two helpers could be made public, but saving the
// match-on-Gamma-enum branch from using them directly (e.g. if one
// knows that the shape is always > 1) doesn't appear to be much
// faster.

/// Gamma distribution where the shape parameter is less than 1.
///
/// Note, samples from this require a compulsory floating-point `pow`
/// call, which makes it significantly slower than sampling from a
/// gamma distribution where the shape parameter is greater than or
/// equal to 1.
///
/// See `Gamma` for sampling from a Gamma distribution with general
/// shape parameters.
struct GammaSmallShape {
    inv_shape: f64,
    large_shape: GammaLargeShape
}

/// Gamma distribution where the shape parameter is larger than 1.
///
/// See `Gamma` for sampling from a Gamma distribution with general
/// shape parameters.
struct GammaLargeShape {
    shape: f64,
    scale: f64,
    c: f64,
    d: f64
}

impl Gamma {
    /// Construct an object representing the `Gamma(shape, scale)`
    /// distribution.
    ///
    /// Fails if `shape <= 0` or `scale <= 0`.
    pub fn new(shape: f64, scale: f64) -> Gamma {
        assert!(shape > 0.0, "Gamma::new called with shape <= 0");
        assert!(scale > 0.0, "Gamma::new called with scale <= 0");

        let repr = match shape {
            1.0        => One(Exp::new(1.0 / scale)),
            0.0 .. 1.0 => Small(GammaSmallShape::new_raw(shape, scale)),
            _          => Large(GammaLargeShape::new_raw(shape, scale))
        };
        Gamma { repr: repr }
    }
}

impl GammaSmallShape {
    fn new_raw(shape: f64, scale: f64) -> GammaSmallShape {
        GammaSmallShape {
            inv_shape: 1. / shape,
            large_shape: GammaLargeShape::new_raw(shape + 1.0, scale)
        }
    }
}

impl GammaLargeShape {
    fn new_raw(shape: f64, scale: f64) -> GammaLargeShape {
        let d = shape - 1. / 3.;
        GammaLargeShape {
            shape: shape,
            scale: scale,
            c: 1. / (9. * d).sqrt(),
            d: d
        }
    }
}

impl Sample<f64> for Gamma {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl Sample<f64> for GammaSmallShape {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl Sample<f64> for GammaLargeShape {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}

impl IndependentSample<f64> for Gamma {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        match self.repr {
            Small(ref g) => g.ind_sample(rng),
            One(ref g) => g.ind_sample(rng),
            Large(ref g) => g.ind_sample(rng),
        }
    }
}
impl IndependentSample<f64> for GammaSmallShape {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        let Open01(u) = rng.gen::<Open01<f64>>();

        self.large_shape.ind_sample(rng) * u.powf(self.inv_shape)
    }
}
impl IndependentSample<f64> for GammaLargeShape {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        loop {
            let StandardNormal(x) = rng.gen::<StandardNormal>();
            let v_cbrt = 1.0 + self.c * x;
            if v_cbrt <= 0.0 { // a^3 <= 0 iff a <= 0
                continue
            }

            let v = v_cbrt * v_cbrt * v_cbrt;
            let Open01(u) = rng.gen::<Open01<f64>>();

            let x_sqr = x * x;
            if u < 1.0 - 0.0331 * x_sqr * x_sqr ||
                u.ln() < 0.5 * x_sqr + self.d * (1.0 - v + v.ln()) {
                return self.d * v * self.scale
            }
        }
    }
}

/// The chi-squared distribution `χ²(k)`, where `k` is the degrees of
/// freedom.
///
/// For `k > 0` integral, this distribution is the sum of the squares
/// of `k` independent standard normal random variables. For other
/// `k`, this uses the equivalent characterisation `χ²(k) = Gamma(k/2,
/// 2)`.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{ChiSquared, IndependentSample};
///
/// let chi = ChiSquared::new(11.0);
/// let v = chi.ind_sample(&mut rand::task_rng());
/// println!("{} is from a χ²(11) distribution", v)
/// ```
pub struct ChiSquared {
    repr: ChiSquaredRepr,
}

enum ChiSquaredRepr {
    // k == 1, Gamma(alpha, ..) is particularly slow for alpha < 1,
    // e.g. when alpha = 1/2 as it would be for this case, so special-
    // casing and using the definition of N(0,1)^2 is faster.
    DoFExactlyOne,
    DoFAnythingElse(Gamma),
}

impl ChiSquared {
    /// Create a new chi-squared distribution with degrees-of-freedom
    /// `k`. Fails if `k < 0`.
    pub fn new(k: f64) -> ChiSquared {
        let repr = if k == 1.0 {
            DoFExactlyOne
        } else {
            assert!(k > 0.0, "ChiSquared::new called with `k` < 0");
            DoFAnythingElse(Gamma::new(0.5 * k, 2.0))
        };
        ChiSquared { repr: repr }
    }
}
impl Sample<f64> for ChiSquared {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl IndependentSample<f64> for ChiSquared {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        match self.repr {
            DoFExactlyOne => {
                // k == 1 => N(0,1)^2
                let StandardNormal(norm) = rng.gen::<StandardNormal>();
                norm * norm
            }
            DoFAnythingElse(ref g) => g.ind_sample(rng)
        }
    }
}

/// The Fisher F distribution `F(m, n)`.
///
/// This distribution is equivalent to the ratio of two normalised
/// chi-squared distributions, that is, `F(m,n) = (χ²(m)/m) /
/// (χ²(n)/n)`.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{FisherF, IndependentSample};
///
/// let f = FisherF::new(2.0, 32.0);
/// let v = f.ind_sample(&mut rand::task_rng());
/// println!("{} is from an F(2, 32) distribution", v)
/// ```
pub struct FisherF {
    numer: ChiSquared,
    denom: ChiSquared,
    // denom_dof / numer_dof so that this can just be a straight
    // multiplication, rather than a division.
    dof_ratio: f64,
}

impl FisherF {
    /// Create a new `FisherF` distribution, with the given
    /// parameter. Fails if either `m` or `n` are not positive.
    pub fn new(m: f64, n: f64) -> FisherF {
        assert!(m > 0.0, "FisherF::new called with `m < 0`");
        assert!(n > 0.0, "FisherF::new called with `n < 0`");

        FisherF {
            numer: ChiSquared::new(m),
            denom: ChiSquared::new(n),
            dof_ratio: n / m
        }
    }
}
impl Sample<f64> for FisherF {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl IndependentSample<f64> for FisherF {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        self.numer.ind_sample(rng) / self.denom.ind_sample(rng) * self.dof_ratio
    }
}

/// The Student t distribution, `t(nu)`, where `nu` is the degrees of
/// freedom.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{StudentT, IndependentSample};
///
/// let t = StudentT::new(11.0);
/// let v = t.ind_sample(&mut rand::task_rng());
/// println!("{} is from a t(11) distribution", v)
/// ```
pub struct StudentT {
    chi: ChiSquared,
    dof: f64
}

impl StudentT {
    /// Create a new Student t distribution with `n` degrees of
    /// freedom. Fails if `n <= 0`.
    pub fn new(n: f64) -> StudentT {
        assert!(n > 0.0, "StudentT::new called with `n <= 0`");
        StudentT {
            chi: ChiSquared::new(n),
            dof: n
        }
    }
}
impl Sample<f64> for StudentT {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl IndependentSample<f64> for StudentT {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        let StandardNormal(norm) = rng.gen::<StandardNormal>();
        norm * (self.dof / self.chi.ind_sample(rng)).sqrt()
    }
}

#[cfg(test)]
mod test {
    use std::prelude::*;

    use distributions::{Sample, IndependentSample};
    use super::{ChiSquared, StudentT, FisherF};

    #[test]
    fn test_chi_squared_one() {
        let mut chi = ChiSquared::new(1.0);
        let mut rng = ::test::rng();
        for _ in range(0, 1000) {
            chi.sample(&mut rng);
            chi.ind_sample(&mut rng);
        }
    }
    #[test]
    fn test_chi_squared_small() {
        let mut chi = ChiSquared::new(0.5);
        let mut rng = ::test::rng();
        for _ in range(0, 1000) {
            chi.sample(&mut rng);
            chi.ind_sample(&mut rng);
        }
    }
    #[test]
    fn test_chi_squared_large() {
        let mut chi = ChiSquared::new(30.0);
        let mut rng = ::test::rng();
        for _ in range(0, 1000) {
            chi.sample(&mut rng);
            chi.ind_sample(&mut rng);
        }
    }
    #[test]
    #[should_fail]
    fn test_chi_squared_invalid_dof() {
        ChiSquared::new(-1.0);
    }

    #[test]
    fn test_f() {
        let mut f = FisherF::new(2.0, 32.0);
        let mut rng = ::test::rng();
        for _ in range(0, 1000) {
            f.sample(&mut rng);
            f.ind_sample(&mut rng);
        }
    }

    #[test]
    fn test_t() {
        let mut t = StudentT::new(11.0);
        let mut rng = ::test::rng();
        for _ in range(0, 1000) {
            t.sample(&mut rng);
            t.ind_sample(&mut rng);
        }
    }
}

#[cfg(test)]
mod bench {
    extern crate test;
    use std::prelude::*;
    use self::test::Bencher;
    use std::mem::size_of;
    use distributions::IndependentSample;
    use super::Gamma;

    #[bench]
    fn bench_gamma_large_shape(b: &mut Bencher) {
        let gamma = Gamma::new(10., 1.0);
        let mut rng = ::test::weak_rng();

        b.iter(|| {
            for _ in range(0, ::RAND_BENCH_N) {
                gamma.ind_sample(&mut rng);
            }
        });
        b.bytes = size_of::<f64>() as u64 * ::RAND_BENCH_N;
    }

    #[bench]
    fn bench_gamma_small_shape(b: &mut Bencher) {
        let gamma = Gamma::new(0.1, 1.0);
        let mut rng = ::test::weak_rng();

        b.iter(|| {
            for _ in range(0, ::RAND_BENCH_N) {
                gamma.ind_sample(&mut rng);
            }
        });
        b.bytes = size_of::<f64>() as u64 * ::RAND_BENCH_N;
    }
}
<|endoftext|>


<commit_before>use rustc_hir::Mutability;
use rustc_index::bit_set::HybridBitSet;
use rustc_middle::mir::visit::{MutVisitor, PlaceContext, Visitor};
use rustc_middle::mir::{self, BasicBlock, Local, Location};
use rustc_middle::ty::TyCtxt;

use crate::transform::{MirPass, MirSource};

/// This pass looks for MIR that always copies the same local into the return place and eliminates
/// the copy by renaming all uses of that local to `_0`.
///
/// This allows LLVM to perform an optimization similar to the named return value optimization
/// (NRVO) that is guaranteed in C++. This avoids a stack allocation and `memcpy` for the
/// relatively common pattern of allocating a buffer on the stack, mutating it, and returning it by
/// value like so:
///
/// ```rust
/// fn foo(init: fn(&mut [u8; 1024])) -> [u8; 1024] {
///     let mut buf = [0; 1024];
///     init(&mut buf);
///     buf
/// }
/// ```
///
/// For now, this pass is very simple and only capable of eliminating a single copy. A more general
/// version of copy propagation, such as the one based on non-overlapping live ranges in [#47954] and
/// [#71003], could yield even more benefits.
///
/// [#47954]: https://github.com/rust-lang/rust/pull/47954
/// [#71003]: https://github.com/rust-lang/rust/pull/71003
pub struct RenameReturnPlace;

impl<'tcx> MirPass<'tcx> for RenameReturnPlace {
    fn run_pass(&self, tcx: TyCtxt<'tcx>, src: MirSource<'tcx>, body: &mut mir::Body<'tcx>) {
        if tcx.sess.opts.debugging_opts.mir_opt_level == 0 {
            return;
        }

        let returned_local = match local_eligible_for_nrvo(body) {
            Some(l) => l,
            None => {
                debug!("`{:?}` was ineligible for NRVO", src.def_id());
                return;
            }
        };

        debug!(
            "`{:?}` was eligible for NRVO, making {:?} the return place",
            src.def_id(),
            returned_local
        );

        RenameToReturnPlace { tcx, to_rename: returned_local }.visit_body(body);

        // Clean up the `NOP`s we inserted for statements made useless by our renaming.
        for block_data in body.basic_blocks_mut() {
            block_data.statements.retain(|stmt| stmt.kind != mir::StatementKind::Nop);
        }

        // Overwrite the debuginfo of `_0` with that of the renamed local.
        let (renamed_decl, ret_decl) =
            body.local_decls.pick2_mut(returned_local, mir::RETURN_PLACE);
        debug_assert_eq!(ret_decl.ty, renamed_decl.ty);
        ret_decl.clone_from(renamed_decl);

        // The return place is always mutable.
        ret_decl.mutability = Mutability::Mut;
    }
}

/// MIR that is eligible for the NRVO must fulfill two conditions:
///   1. The return place must not be read prior to the `Return` terminator.
///   2. A simple assignment of a whole local to the return place (e.g., `_0 = _1`) must be the
///      only definition of the return place reaching the `Return` terminator.
///
/// If the MIR fulfills both these conditions, this function returns the `Local` that is assigned
/// to the return place along all possible paths through the control-flow graph.
fn local_eligible_for_nrvo(body: &mut mir::Body<'_>) -> Option<Local> {
    if IsReturnPlaceRead::run(body) {
        return None;
    }

    let mut copied_to_return_place = None;
    for block in body.basic_blocks().indices() {
        // Look for blocks with a `Return` terminator.
        if !matches!(body[block].terminator().kind, mir::TerminatorKind::Return) {
            continue;
        }

        // Look for an assignment of a single local to the return place prior to the `Return`.
        let returned_local = find_local_assigned_to_return_place(block, body)?;
        match body.local_kind(returned_local) {
            // FIXME: Can we do this for arguments as well?
            mir::LocalKind::Arg => return None,

            mir::LocalKind::ReturnPointer => bug!("Return place was assigned to itself?"),
            mir::LocalKind::Var | mir::LocalKind::Temp => {}
        }

        // If multiple different locals are copied to the return place. We can't pick a
        // single one to rename.
        if copied_to_return_place.map_or(false, |old| old != returned_local) {
            return None;
        }

        copied_to_return_place = Some(returned_local);
    }

    return copied_to_return_place;
}

fn find_local_assigned_to_return_place(
    start: BasicBlock,
    body: &mut mir::Body<'_>,
) -> Option<Local> {
    let mut block = start;
    let mut seen = HybridBitSet::new_empty(body.basic_blocks().len());

    // Iterate as long as `block` has exactly one predecessor that we have not yet visited.
    while seen.insert(block) {
        trace!("Looking for assignments to `_0` in {:?}", block);

        let local = body[block].statements.iter().rev().find_map(as_local_assigned_to_return_place);
        if local.is_some() {
            return local;
        }

        match body.predecessors()[block].as_slice() {
            &[pred] => block = pred,
            _ => return None,
        }
    }

    return None;
}

// If this statement is an assignment of an unprojected local to the return place,
// return that local.
fn as_local_assigned_to_return_place(stmt: &mir::Statement<'_>) -> Option<Local> {
    if let mir::StatementKind::Assign(box (lhs, rhs)) = &stmt.kind {
        if lhs.as_local() == Some(mir::RETURN_PLACE) {
            if let mir::Rvalue::Use(mir::Operand::Copy(rhs) | mir::Operand::Move(rhs)) = rhs {
                return rhs.as_local();
            }
        }
    }

    None
}

struct RenameToReturnPlace<'tcx> {
    to_rename: Local,
    tcx: TyCtxt<'tcx>,
}

/// Replaces all uses of `self.to_rename` with `_0`.
impl MutVisitor<'tcx> for RenameToReturnPlace<'tcx> {
    fn tcx(&self) -> TyCtxt<'tcx> {
        self.tcx
    }

    fn visit_statement(&mut self, stmt: &mut mir::Statement<'tcx>, loc: Location) {
        // Remove assignments of the local being replaced to the return place, since it is now the
        // return place:
        //     _0 = _1
        if as_local_assigned_to_return_place(stmt) == Some(self.to_rename) {
            stmt.kind = mir::StatementKind::Nop;
            return;
        }

        // Remove storage annotations for the local being replaced:
        //     StorageLive(_1)
        if let mir::StatementKind::StorageLive(local) | mir::StatementKind::StorageDead(local) =
            stmt.kind
        {
            if local == self.to_rename {
                stmt.kind = mir::StatementKind::Nop;
                return;
            }
        }

        self.super_statement(stmt, loc)
    }

    fn visit_terminator(&mut self, terminator: &mut mir::Terminator<'tcx>, loc: Location) {
        // Ignore the implicit "use" of the return place in a `Return` statement.
        if let mir::TerminatorKind::Return = terminator.kind {
            return;
        }

        self.super_terminator(terminator, loc);
    }

    fn visit_local(&mut self, l: &mut Local, _: PlaceContext, _: Location) {
        assert_ne!(*l, mir::RETURN_PLACE);
        if *l == self.to_rename {
            *l = mir::RETURN_PLACE;
        }
    }
}

struct IsReturnPlaceRead(bool);

impl IsReturnPlaceRead {
    fn run(body: &mir::Body<'_>) -> bool {
        let mut vis = IsReturnPlaceRead(false);
        vis.visit_body(body);
        vis.0
    }
}

impl Visitor<'tcx> for IsReturnPlaceRead {
    fn visit_local(&mut self, &l: &Local, ctxt: PlaceContext, _: Location) {
        if l == mir::RETURN_PLACE && ctxt.is_use() && !ctxt.is_place_assignment() {
            self.0 = true;
        }
    }

    fn visit_terminator(&mut self, terminator: &mir::Terminator<'tcx>, loc: Location) {
        // Ignore the implicit "use" of the return place in a `Return` statement.
        if let mir::TerminatorKind::Return = terminator.kind {
            return;
        }

        self.super_terminator(terminator, loc);
    }
}
<commit_msg>Bail out if new return place has different type than old<commit_after>use rustc_hir::Mutability;
use rustc_index::bit_set::HybridBitSet;
use rustc_middle::mir::visit::{MutVisitor, PlaceContext, Visitor};
use rustc_middle::mir::{self, BasicBlock, Local, Location};
use rustc_middle::ty::TyCtxt;

use crate::transform::{MirPass, MirSource};

/// This pass looks for MIR that always copies the same local into the return place and eliminates
/// the copy by renaming all uses of that local to `_0`.
///
/// This allows LLVM to perform an optimization similar to the named return value optimization
/// (NRVO) that is guaranteed in C++. This avoids a stack allocation and `memcpy` for the
/// relatively common pattern of allocating a buffer on the stack, mutating it, and returning it by
/// value like so:
///
/// ```rust
/// fn foo(init: fn(&mut [u8; 1024])) -> [u8; 1024] {
///     let mut buf = [0; 1024];
///     init(&mut buf);
///     buf
/// }
/// ```
///
/// For now, this pass is very simple and only capable of eliminating a single copy. A more general
/// version of copy propagation, such as the one based on non-overlapping live ranges in [#47954] and
/// [#71003], could yield even more benefits.
///
/// [#47954]: https://github.com/rust-lang/rust/pull/47954
/// [#71003]: https://github.com/rust-lang/rust/pull/71003
pub struct RenameReturnPlace;

impl<'tcx> MirPass<'tcx> for RenameReturnPlace {
    fn run_pass(&self, tcx: TyCtxt<'tcx>, src: MirSource<'tcx>, body: &mut mir::Body<'tcx>) {
        if tcx.sess.opts.debugging_opts.mir_opt_level == 0 {
            return;
        }

        let returned_local = match local_eligible_for_nrvo(body) {
            Some(l) => l,
            None => {
                debug!("`{:?}` was ineligible for NRVO", src.def_id());
                return;
            }
        };

        // Sometimes, the return place is assigned a local of a different but coercable type, for
        // example `&T` instead of `&mut T`. Overwriting the `LocalInfo` for the return place would
        // result in it having an incorrect type. Although this doesn't seem to cause a problem in
        // codegen, bail out anyways since it happens so rarely.
        let ret_ty = body.local_decls[mir::RETURN_PLACE].ty;
        let assigned_ty = body.local_decls[returned_local].ty;
        if ret_ty != assigned_ty {
            debug!("`{:?}` was eligible for NRVO but for type mismatch", src.def_id());
            debug!("typeof(_0) != typeof({:?}); {:?} != {:?}", returned_local, ret_ty, assigned_ty);
            return;
        }

        debug!(
            "`{:?}` was eligible for NRVO, making {:?} the return place",
            src.def_id(),
            returned_local
        );

        RenameToReturnPlace { tcx, to_rename: returned_local }.visit_body(body);

        // Clean up the `NOP`s we inserted for statements made useless by our renaming.
        for block_data in body.basic_blocks_mut() {
            block_data.statements.retain(|stmt| stmt.kind != mir::StatementKind::Nop);
        }

        // Overwrite the debuginfo of `_0` with that of the renamed local.
        let (renamed_decl, ret_decl) =
            body.local_decls.pick2_mut(returned_local, mir::RETURN_PLACE);
        ret_decl.clone_from(renamed_decl);

        // The return place is always mutable.
        ret_decl.mutability = Mutability::Mut;
    }
}

/// MIR that is eligible for the NRVO must fulfill two conditions:
///   1. The return place must not be read prior to the `Return` terminator.
///   2. A simple assignment of a whole local to the return place (e.g., `_0 = _1`) must be the
///      only definition of the return place reaching the `Return` terminator.
///
/// If the MIR fulfills both these conditions, this function returns the `Local` that is assigned
/// to the return place along all possible paths through the control-flow graph.
fn local_eligible_for_nrvo(body: &mut mir::Body<'_>) -> Option<Local> {
    if IsReturnPlaceRead::run(body) {
        return None;
    }

    let mut copied_to_return_place = None;
    for block in body.basic_blocks().indices() {
        // Look for blocks with a `Return` terminator.
        if !matches!(body[block].terminator().kind, mir::TerminatorKind::Return) {
            continue;
        }

        // Look for an assignment of a single local to the return place prior to the `Return`.
        let returned_local = find_local_assigned_to_return_place(block, body)?;
        match body.local_kind(returned_local) {
            // FIXME: Can we do this for arguments as well?
            mir::LocalKind::Arg => return None,

            mir::LocalKind::ReturnPointer => bug!("Return place was assigned to itself?"),
            mir::LocalKind::Var | mir::LocalKind::Temp => {}
        }

        // If multiple different locals are copied to the return place. We can't pick a
        // single one to rename.
        if copied_to_return_place.map_or(false, |old| old != returned_local) {
            return None;
        }

        copied_to_return_place = Some(returned_local);
    }

    return copied_to_return_place;
}

fn find_local_assigned_to_return_place(
    start: BasicBlock,
    body: &mut mir::Body<'_>,
) -> Option<Local> {
    let mut block = start;
    let mut seen = HybridBitSet::new_empty(body.basic_blocks().len());

    // Iterate as long as `block` has exactly one predecessor that we have not yet visited.
    while seen.insert(block) {
        trace!("Looking for assignments to `_0` in {:?}", block);

        let local = body[block].statements.iter().rev().find_map(as_local_assigned_to_return_place);
        if local.is_some() {
            return local;
        }

        match body.predecessors()[block].as_slice() {
            &[pred] => block = pred,
            _ => return None,
        }
    }

    return None;
}

// If this statement is an assignment of an unprojected local to the return place,
// return that local.
fn as_local_assigned_to_return_place(stmt: &mir::Statement<'_>) -> Option<Local> {
    if let mir::StatementKind::Assign(box (lhs, rhs)) = &stmt.kind {
        if lhs.as_local() == Some(mir::RETURN_PLACE) {
            if let mir::Rvalue::Use(mir::Operand::Copy(rhs) | mir::Operand::Move(rhs)) = rhs {
                return rhs.as_local();
            }
        }
    }

    None
}

struct RenameToReturnPlace<'tcx> {
    to_rename: Local,
    tcx: TyCtxt<'tcx>,
}

/// Replaces all uses of `self.to_rename` with `_0`.
impl MutVisitor<'tcx> for RenameToReturnPlace<'tcx> {
    fn tcx(&self) -> TyCtxt<'tcx> {
        self.tcx
    }

    fn visit_statement(&mut self, stmt: &mut mir::Statement<'tcx>, loc: Location) {
        // Remove assignments of the local being replaced to the return place, since it is now the
        // return place:
        //     _0 = _1
        if as_local_assigned_to_return_place(stmt) == Some(self.to_rename) {
            stmt.kind = mir::StatementKind::Nop;
            return;
        }

        // Remove storage annotations for the local being replaced:
        //     StorageLive(_1)
        if let mir::StatementKind::StorageLive(local) | mir::StatementKind::StorageDead(local) =
            stmt.kind
        {
            if local == self.to_rename {
                stmt.kind = mir::StatementKind::Nop;
                return;
            }
        }

        self.super_statement(stmt, loc)
    }

    fn visit_terminator(&mut self, terminator: &mut mir::Terminator<'tcx>, loc: Location) {
        // Ignore the implicit "use" of the return place in a `Return` statement.
        if let mir::TerminatorKind::Return = terminator.kind {
            return;
        }

        self.super_terminator(terminator, loc);
    }

    fn visit_local(&mut self, l: &mut Local, _: PlaceContext, _: Location) {
        assert_ne!(*l, mir::RETURN_PLACE);
        if *l == self.to_rename {
            *l = mir::RETURN_PLACE;
        }
    }
}

struct IsReturnPlaceRead(bool);

impl IsReturnPlaceRead {
    fn run(body: &mir::Body<'_>) -> bool {
        let mut vis = IsReturnPlaceRead(false);
        vis.visit_body(body);
        vis.0
    }
}

impl Visitor<'tcx> for IsReturnPlaceRead {
    fn visit_local(&mut self, &l: &Local, ctxt: PlaceContext, _: Location) {
        if l == mir::RETURN_PLACE && ctxt.is_use() && !ctxt.is_place_assignment() {
            self.0 = true;
        }
    }

    fn visit_terminator(&mut self, terminator: &mir::Terminator<'tcx>, loc: Location) {
        // Ignore the implicit "use" of the return place in a `Return` statement.
        if let mir::TerminatorKind::Return = terminator.kind {
            return;
        }

        self.super_terminator(terminator, loc);
    }
}
<|endoftext|>


<commit_before>// compile-flags: -Zmiri-seed=0000000000000000

fn main() {
    // Some casting-to-int with arithmetic.
    let x = &42 as *const i32 as usize;
    let y = x * 2;
    assert_eq!(y, x + x);
    let z = y as u8 as usize;
    assert_eq!(z, y % 256);

    // Pointer string formatting! We can't check the output as it changes when libstd changes,
    // but we can make sure Miri does not error.
    format!("{:?}", &mut 13 as *mut _);
}
<commit_msg>Add intptrcast test for explicit casts<commit_after>// compile-flags: -Zmiri-seed=0000000000000000
fn transmute_ptr_to_int<T>(x: *const T) -> usize {
    unsafe { std::mem::transmute::<*const T, usize>(x) * 1 }
}

fn main() {
    // Some casting-to-int with arithmetic.
    let x = &42 as *const i32 as usize;
    let y = x * 2;
    assert_eq!(y, x + x);
    let z = y as u8 as usize;
    assert_eq!(z, y % 256);

    // Pointer string formatting! We can't check the output as it changes when libstd changes,
    // but we can make sure Miri does not error.
    format!("{:?}", &mut 13 as *mut _);

    // Check that intptrcast is triggered for explicit casts and that it is consistent with
    // transmuting.
    let a: *const i32 = &42;
    let b = transmute_ptr_to_int(a) as u8;
    let c = a as usize as u8;
    assert_eq!(b, c);
}
<|endoftext|>


<commit_before>use std::cell::RefCell;
use std::cmp;

use rand;
use rand::Rng;
use itertools::Itertools;
use rgsl::randist::multinomial::multinomial_pdf;

use bio::stats::{Prob, LogProb};

use io::codebook::{Codebook, FeatureID};

use model::readout::{Expressions, Counts, Miscalls, Xi};

pub trait AbstractFeatureModel {
    /// FeatureID of this feature.
    fn feature_id(&self) -> FeatureID;

    /// Neighbors of this feature.
    fn neighbors(&self) -> &[FeatureID];

    /// Minimum hamming distance between any pair of features.
    fn min_dist(&self) -> u8;

    /// Number of events (calls and miscalls) except dropout.
    fn event_count_non_dropout(
        &self,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> u32;

    /// Dropout probability.
    fn prob_dropout(&self) -> f64;

    /// Probability of exact miscall to given neighbor index.
    fn prob_miscall_exact(&self, neighbor_index: usize) -> f64;

    /// Probability of mismatching miscall to given neighbor index.
    fn prob_miscall_mismatch(&self, neighbor_index: usize) -> f64;

    /// Calculate maximum likelihood estimate of miscalls given expression and store in the given
    /// arrays.
    /// Returns the total absolute changes to estimated miscall counts.
    fn mle_miscalls(
        &self,
        expressions: &Expressions,
        miscalls_exact: &mut Miscalls,
        miscalls_mismatch: &mut Miscalls,
        rng: &mut rand::StdRng
    ) {
        let x = expressions[self.feature_id()];

        // We need to ensure that the sum of the estimated miscalls does not exceed x.
        // Hence, we iterate randomly over the neighbors such that we can simply stop once that
        // happens without causing a bias.
        let mut idx = (0..self.neighbors().len()).collect_vec();
        rng.shuffle(&mut idx);

        // by stochastic rounding, we ensure that on average \sum x*p_i = x
        // numeric rounding would not work if many of the probabilities are small enough to yield
        // expectations less than 1
        let mut stochastic_round = |v: f64| v.floor() as u32 + (rng.next_f64() <= v % 1.0) as u32;

        let mut rest = x;
        for i in idx {
            let n = self.neighbors()[i];
            let miscall_exact = self.prob_miscall_exact(i) * x as f64;
            miscalls_exact.set(
                self.feature_id(), n, stochastic_round(miscall_exact), &mut rest
            );

            if self.min_dist() == 4 {
                let miscall_mismatch = self.prob_miscall_mismatch(i) * x as f64;
                let miscall_mismatch = stochastic_round(miscall_mismatch);
                miscalls_mismatch.set(
                    self.feature_id(), n, miscall_mismatch, &mut rest
                );
            }
        }
    }

    /// Calculate maximum likelihood estimate of expression given miscalls.
    /// Returns the absolute amount of change to the expression.
    fn mle_expression(
        &self,
        expressions: &mut Expressions,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> i32 {
        let event_count = self.event_count_non_dropout(miscalls_exact, miscalls_mismatch);

        let x_ = cmp::max(
            (event_count as f64 / (1.0 - self.prob_dropout())).floor() as u32,
            // or at least as many as we have events coming from this feature
            // TODO does this allow x to decrease properly?
            event_count
        );

        let x = expressions.get_mut(self.feature_id()).unwrap();

        let change = *x as i32 - x_ as i32;

        *x = x_;

        change
    }
}

pub struct FeatureModel {
    pub(crate) feature_id: FeatureID,
    event_probs: Vec<f64>,
    neighbors: Vec<FeatureID>,
    min_dist: u8,
    pub(crate) counts: Counts,
    event_counts: RefCell<Vec<u32>>
}

impl FeatureModel {
    pub fn new(
        feature: FeatureID,
        counts: Counts,
        codebook: &Codebook,
        xi: &Xi
    ) -> Self {
        assert!(
            codebook.min_dist == 2 || codebook.min_dist == 4,
            "unsupported hamming distance: {}, supported: 2, 4",
            codebook.min_dist
        );

        let feature_record = codebook.record(feature);

        let xi_call = xi.prob(feature_record.codeword(), feature_record.codeword());

        let prob_call_exact = xi_call[0];
        let prob_call_mismatch = if codebook.min_dist == 4 {
            xi_call[1]
        } else {
            LogProb::ln_zero()
        };

        let neighbors_4 = codebook.neighbors(feature, 4);
        let neighbors_2 = if codebook.min_dist == 2 {
            codebook.neighbors(feature, 2)
        } else {
            vec![]
        };

        let mut neighbors = neighbors_4.clone();
        neighbors.extend(&neighbors_2);

        let xi_miscall_4 = codebook.neighbors(feature, 4).iter().map(|&n| {
            xi.prob(feature_record.codeword(), codebook.record(n).codeword())
        }).collect_vec();
        let xi_miscall_2 = if codebook.min_dist == 2 {
            codebook.neighbors(feature, 2).iter().map(|&n| {
                xi.prob(feature_record.codeword(), codebook.record(n).codeword())
            }).collect_vec()
        } else {
            vec![]
        };

        let prob_miscall_exact = xi_miscall_4.iter().map(|xi| {
            xi[0]
        }).chain(xi_miscall_2.iter().map(|xi| xi[0])).collect_vec();

        let prob_miscall_mismatch = if codebook.min_dist == 4 {
            xi_miscall_4.iter().map(|xi| xi[1]).collect_vec()
        } else {
            vec![]
        };

        let mut total = vec![prob_call_exact, prob_call_mismatch];
        total.extend(&prob_miscall_exact);
        total.extend(&prob_miscall_mismatch);

        let prob_total = LogProb::ln_sum_exp(&total);
        assert!(*prob_total <= 0.0, "prob_total = {} > 0", *prob_total);

        let mut event_probs = vec![
            *Prob::from(prob_call_exact),
            *Prob::from(prob_call_mismatch),
            *Prob::from(prob_total.ln_one_minus_exp())
        ];
        event_probs.extend(prob_miscall_exact.iter().map(|&p| *Prob::from(p)));
        event_probs.extend(prob_miscall_mismatch.iter().map(|&p| *Prob::from(p)));

        FeatureModel {
            feature_id: feature,
            event_probs: event_probs,
            neighbors: neighbors,
            min_dist: codebook.min_dist,
            counts: counts,
            event_counts: RefCell::new(Vec::new())
        }
    }

    /// Calculate likelihood of given expression.
    pub fn likelihood(
        &self,
        x: u32,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> LogProb {
        let count_exact_miscalled = miscalls_exact.total_to(self.feature_id);
        let count_mismatch_miscalled = miscalls_mismatch.total_to(self.feature_id);

        if self.counts.exact < count_exact_miscalled {
            panic!("bug: likelihood cannot be calculated if exact miscalls exceed counts: counts={}, miscalls={}", self.counts.exact, count_exact_miscalled);
        }
        if self.counts.mismatch < count_mismatch_miscalled {
            panic!("bug: likelihood cannot be calculated if mismatch miscalls exceed counts: counts={}, miscalls={}", self.counts.mismatch, count_mismatch_miscalled);
        }
        let total_count = self.counts.exact - count_exact_miscalled +
                          self.counts.mismatch - count_mismatch_miscalled;
        if total_count > x {
            panic!("bug: event count exceeds expression: x={}, event count={}", x, total_count);
        }

        let calls_exact = self.counts.exact - count_exact_miscalled;
        let calls_mismatch = self.counts.mismatch - count_mismatch_miscalled;

        // setup event counts
        {
            let mut event_counts = self.event_counts.borrow_mut();
            event_counts.clear();
            event_counts.push(
                // exact calls
                calls_exact
            );
            event_counts.push(
                // mismatch calls
                calls_mismatch
            );
            event_counts.push(
                // dropouts (fill in later)
                0
            );
            // miscalls
            event_counts.extend(self.neighbors.iter().map(
                |&n| miscalls_exact.get(self.feature_id, n)
            ));
            if self.min_dist == 4 {
                event_counts.extend(
                    self.neighbors.iter().map(
                        |&n| miscalls_mismatch.get(self.feature_id, n)
                    )
                );
            }

            // set dropouts
            let total_counts = event_counts.iter().sum();
            assert!(x >= total_counts, "total event counts exceed expression: x={} {:?}", x, event_counts);
            let dropouts = x - total_counts;
            event_counts[2] = dropouts;
        }

        LogProb::from(Prob(
            multinomial_pdf(&self.event_probs, &self.event_counts.borrow())
        ))
    }

    /// Minimum expression given fixed miscalls.
    pub fn min_expression(&self, miscalls_exact: &Miscalls, miscalls_mismatch: &Miscalls) -> u32 {
        self.counts.exact.saturating_sub(miscalls_exact.total_to(self.feature_id)) +
        self.counts.mismatch.saturating_sub(miscalls_mismatch.total_to(self.feature_id)) +
        miscalls_exact.total_from(self.feature_id) + miscalls_mismatch.total_from(self.feature_id)
    }
}

impl AbstractFeatureModel for FeatureModel {
    fn feature_id(&self) -> FeatureID {
        self.feature_id
    }

    fn neighbors(&self) -> &[FeatureID] {
        &self.neighbors
    }

    fn min_dist(&self) -> u8 {
        self.min_dist
    }

    fn event_count_non_dropout(
        &self,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> u32 {
        let calls_exact = self.counts.exact.saturating_sub(
            miscalls_exact.total_to(self.feature_id)
        );
        let calls_mismatch = self.counts.mismatch.saturating_sub(
            miscalls_mismatch.total_to(self.feature_id)
        );

        calls_exact + calls_mismatch +
        miscalls_exact.total_from(self.feature_id) +
        miscalls_mismatch.total_from(self.feature_id)
    }

    fn prob_dropout(&self) -> f64 {
        self.event_probs[2]
    }

    fn prob_miscall_exact(&self, neighbor_index: usize) -> f64 {
        self.event_probs[3 + neighbor_index]
    }

    fn prob_miscall_mismatch(&self, neighbor_index: usize) -> f64 {
        self.event_probs[3 + self.neighbors.len() + neighbor_index]
    }
}

pub struct NoiseModel {
    pub(crate) feature_id: FeatureID,
    pub(crate) not_expressed_feature_ids: Vec<FeatureID>,
    pub(crate) not_expressed_counts: Vec<Counts>,
    probs_miscall_exact: Vec<Prob>,
    probs_miscall_mismatch: Vec<Prob>,
    prob_dropout: Prob,
    neighbors: Vec<FeatureID>,
    min_dist: u8,
}

impl NoiseModel {
    pub fn new(
        feature_id: FeatureID,
        not_expressed_feature_ids: Vec<FeatureID>,
        not_expressed_counts: Vec<Counts>,
        codebook: &Codebook,
        xi: &Xi
    ) -> Self {
        assert!(
            codebook.min_dist == 2 || codebook.min_dist == 4,
            "unsupported hamming distance: {}, supported: 2, 4",
            codebook.min_dist
        );

        // miscalls generated from noise
        let xi_miscall = codebook.records().iter().filter_map(|feat_rec| {
            if feat_rec.expressed() {
                Some(xi.prob(codebook.noise().codeword(), feat_rec.codeword()))
            } else { None }
        }).collect_vec();
        let neighbors = codebook.records().iter().filter_map(|rec| {
            if rec.expressed() {
                Some(codebook.get_id(rec.name()))
            } else { None }
        }).collect_vec();

        let probs_miscall_exact = xi_miscall.iter().map(|xi| xi[0]).collect_vec();
        let probs_miscall_mismatch = if codebook.min_dist == 4 {
            xi_miscall.iter().map(|xi| xi[1]).collect_vec()
        } else {
            vec![]
        };

        let prob_total_miscall_exact = LogProb::ln_sum_exp(&probs_miscall_exact);
        let prob_total_miscall_mismatch = LogProb::ln_sum_exp(&probs_miscall_mismatch);

        // probs for calls of not expressed features
        let xi_not_expressed = not_expressed_feature_ids.iter().map(|&f| {
            xi.prob(codebook.noise().codeword(), codebook.record(f).codeword())
        }).collect_vec();

        let prob_not_expressed_exact = LogProb::ln_sum_exp(
            &xi_not_expressed.iter().map(|xi| xi[0]).collect_vec()
        );
        let prob_not_expressed_mismatch = LogProb::ln_sum_exp(
            &xi_not_expressed.iter().map(|xi| xi[1]).collect_vec()
        );

        // dropout
        let prob_dropout = LogProb::ln_sum_exp(&[
            prob_not_expressed_exact,
            prob_not_expressed_mismatch,
            prob_total_miscall_exact,
            prob_total_miscall_mismatch
        ]).ln_one_minus_exp();

        NoiseModel {
            feature_id: feature_id,
            not_expressed_feature_ids: not_expressed_feature_ids,
            not_expressed_counts: not_expressed_counts,
            probs_miscall_exact: probs_miscall_exact.iter().map(|&p| Prob::from(p)).collect_vec(),
            probs_miscall_mismatch: probs_miscall_mismatch.iter().map(|&p| Prob::from(p)).collect_vec(),
            prob_dropout: Prob::from(prob_dropout),
            neighbors: neighbors,
            min_dist: codebook.min_dist
        }
    }
}

impl AbstractFeatureModel for NoiseModel {
    fn feature_id(&self) -> FeatureID {
        self.feature_id
    }

    fn neighbors(&self) -> &[FeatureID] {
        &self.neighbors
    }

    fn min_dist(&self) -> u8 {
        self.min_dist
    }

    fn event_count_non_dropout(
        &self,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> u32 {
        let mut calls_exact = 0;
        let mut calls_mismatch = 0;
        for (&id, counts) in self.not_expressed_feature_ids.iter().zip(&self.not_expressed_counts) {
            calls_exact += counts.exact.saturating_sub(miscalls_exact.total_to(id));
            calls_mismatch += counts.mismatch.saturating_sub(miscalls_mismatch.total_to(id));
        }

        let event_count = calls_exact + calls_mismatch +
                          miscalls_exact.total_from(self.feature_id) +
                          miscalls_mismatch.total_from(self.feature_id);

        event_count
    }

    fn prob_dropout(&self) -> f64 {
        *self.prob_dropout
    }

    fn prob_miscall_exact(&self, neighbor_index: usize) -> f64 {
        *self.probs_miscall_exact[neighbor_index]
    }

    fn prob_miscall_mismatch(&self, neighbor_index: usize) -> f64 {
        *self.probs_miscall_mismatch[neighbor_index]
    }
}
<commit_msg>Temporarily disable noise.<commit_after>use std::cell::RefCell;
use std::cmp;

use rand;
use rand::Rng;
use itertools::Itertools;
use rgsl::randist::multinomial::multinomial_pdf;

use bio::stats::{Prob, LogProb};

use io::codebook::{Codebook, FeatureID};

use model::readout::{Expressions, Counts, Miscalls, Xi};

pub trait AbstractFeatureModel {
    /// FeatureID of this feature.
    fn feature_id(&self) -> FeatureID;

    /// Neighbors of this feature.
    fn neighbors(&self) -> &[FeatureID];

    /// Minimum hamming distance between any pair of features.
    fn min_dist(&self) -> u8;

    /// Number of events (calls and miscalls) except dropout.
    fn event_count_non_dropout(
        &self,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> u32;

    /// Dropout probability.
    fn prob_dropout(&self) -> f64;

    /// Probability of exact miscall to given neighbor index.
    fn prob_miscall_exact(&self, neighbor_index: usize) -> f64;

    /// Probability of mismatching miscall to given neighbor index.
    fn prob_miscall_mismatch(&self, neighbor_index: usize) -> f64;

    /// Calculate maximum likelihood estimate of miscalls given expression and store in the given
    /// arrays.
    /// Returns the total absolute changes to estimated miscall counts.
    fn mle_miscalls(
        &self,
        expressions: &Expressions,
        miscalls_exact: &mut Miscalls,
        miscalls_mismatch: &mut Miscalls,
        rng: &mut rand::StdRng
    ) {
        let x = expressions[self.feature_id()];

        // We need to ensure that the sum of the estimated miscalls does not exceed x.
        // Hence, we iterate randomly over the neighbors such that we can simply stop once that
        // happens without causing a bias.
        let mut idx = (0..self.neighbors().len()).collect_vec();
        rng.shuffle(&mut idx);

        // by stochastic rounding, we ensure that on average \sum x*p_i = x
        // numeric rounding would not work if many of the probabilities are small enough to yield
        // expectations less than 1
        let mut stochastic_round = |v: f64| v.floor() as u32 + (rng.next_f64() <= v % 1.0) as u32;

        let mut rest = x;
        for i in idx {
            let n = self.neighbors()[i];
            let miscall_exact = self.prob_miscall_exact(i) * x as f64;
            miscalls_exact.set(
                self.feature_id(), n, stochastic_round(miscall_exact), &mut rest
            );

            if self.min_dist() == 4 {
                let miscall_mismatch = self.prob_miscall_mismatch(i) * x as f64;
                let miscall_mismatch = stochastic_round(miscall_mismatch);
                miscalls_mismatch.set(
                    self.feature_id(), n, miscall_mismatch, &mut rest
                );
            }
        }
    }

    /// Calculate maximum likelihood estimate of expression given miscalls.
    /// Returns the absolute amount of change to the expression.
    fn mle_expression(
        &self,
        expressions: &mut Expressions,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> i32 {
        let event_count = self.event_count_non_dropout(miscalls_exact, miscalls_mismatch);

        let x_ = cmp::max(
            (event_count as f64 / (1.0 - self.prob_dropout())).floor() as u32,
            // or at least as many as we have events coming from this feature
            // TODO does this allow x to decrease properly?
            event_count
        );

        let x = expressions.get_mut(self.feature_id()).unwrap();

        let change = *x as i32 - x_ as i32;

        *x = x_;

        change
    }
}

pub struct FeatureModel {
    pub(crate) feature_id: FeatureID,
    event_probs: Vec<f64>,
    neighbors: Vec<FeatureID>,
    min_dist: u8,
    pub(crate) counts: Counts,
    event_counts: RefCell<Vec<u32>>
}

impl FeatureModel {
    pub fn new(
        feature: FeatureID,
        counts: Counts,
        codebook: &Codebook,
        xi: &Xi
    ) -> Self {
        assert!(
            codebook.min_dist == 2 || codebook.min_dist == 4,
            "unsupported hamming distance: {}, supported: 2, 4",
            codebook.min_dist
        );

        let feature_record = codebook.record(feature);

        let xi_call = xi.prob(feature_record.codeword(), feature_record.codeword());

        let prob_call_exact = xi_call[0];
        let prob_call_mismatch = if codebook.min_dist == 4 {
            xi_call[1]
        } else {
            LogProb::ln_zero()
        };

        let neighbors_4 = codebook.neighbors(feature, 4);
        let neighbors_2 = if codebook.min_dist == 2 {
            codebook.neighbors(feature, 2)
        } else {
            vec![]
        };

        let mut neighbors = neighbors_4.clone();
        neighbors.extend(&neighbors_2);

        let xi_miscall_4 = codebook.neighbors(feature, 4).iter().map(|&n| {
            xi.prob(feature_record.codeword(), codebook.record(n).codeword())
        }).collect_vec();
        let xi_miscall_2 = if codebook.min_dist == 2 {
            codebook.neighbors(feature, 2).iter().map(|&n| {
                xi.prob(feature_record.codeword(), codebook.record(n).codeword())
            }).collect_vec()
        } else {
            vec![]
        };

        let prob_miscall_exact = xi_miscall_4.iter().map(|xi| {
            xi[0]
        }).chain(xi_miscall_2.iter().map(|xi| xi[0])).collect_vec();

        let prob_miscall_mismatch = if codebook.min_dist == 4 {
            xi_miscall_4.iter().map(|xi| xi[1]).collect_vec()
        } else {
            vec![]
        };

        let mut total = vec![prob_call_exact, prob_call_mismatch];
        total.extend(&prob_miscall_exact);
        total.extend(&prob_miscall_mismatch);

        let prob_total = LogProb::ln_sum_exp(&total);
        assert!(*prob_total <= 0.0, "prob_total = {} > 0", *prob_total);

        let mut event_probs = vec![
            *Prob::from(prob_call_exact),
            *Prob::from(prob_call_mismatch),
            *Prob::from(prob_total.ln_one_minus_exp())
        ];
        event_probs.extend(prob_miscall_exact.iter().map(|&p| *Prob::from(p)));
        event_probs.extend(prob_miscall_mismatch.iter().map(|&p| *Prob::from(p)));

        FeatureModel {
            feature_id: feature,
            event_probs: event_probs,
            neighbors: neighbors,
            min_dist: codebook.min_dist,
            counts: counts,
            event_counts: RefCell::new(Vec::new())
        }
    }

    /// Calculate likelihood of given expression.
    pub fn likelihood(
        &self,
        x: u32,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> LogProb {
        let count_exact_miscalled = miscalls_exact.total_to(self.feature_id);
        let count_mismatch_miscalled = miscalls_mismatch.total_to(self.feature_id);

        if self.counts.exact < count_exact_miscalled {
            panic!("bug: likelihood cannot be calculated if exact miscalls exceed counts: counts={}, miscalls={}", self.counts.exact, count_exact_miscalled);
        }
        if self.counts.mismatch < count_mismatch_miscalled {
            panic!("bug: likelihood cannot be calculated if mismatch miscalls exceed counts: counts={}, miscalls={}", self.counts.mismatch, count_mismatch_miscalled);
        }
        let total_count = self.counts.exact - count_exact_miscalled +
                          self.counts.mismatch - count_mismatch_miscalled;
        if total_count > x {
            panic!("bug: event count exceeds expression: x={}, event count={}", x, total_count);
        }

        let calls_exact = self.counts.exact - count_exact_miscalled;
        let calls_mismatch = self.counts.mismatch - count_mismatch_miscalled;

        // setup event counts
        {
            let mut event_counts = self.event_counts.borrow_mut();
            event_counts.clear();
            event_counts.push(
                // exact calls
                calls_exact
            );
            event_counts.push(
                // mismatch calls
                calls_mismatch
            );
            event_counts.push(
                // dropouts (fill in later)
                0
            );
            // miscalls
            event_counts.extend(self.neighbors.iter().map(
                |&n| miscalls_exact.get(self.feature_id, n)
            ));
            if self.min_dist == 4 {
                event_counts.extend(
                    self.neighbors.iter().map(
                        |&n| miscalls_mismatch.get(self.feature_id, n)
                    )
                );
            }

            // set dropouts
            let total_counts = event_counts.iter().sum();
            assert!(x >= total_counts, "total event counts exceed expression: x={} {:?}", x, event_counts);
            let dropouts = x - total_counts;
            event_counts[2] = dropouts;
        }

        LogProb::from(Prob(
            multinomial_pdf(&self.event_probs, &self.event_counts.borrow())
        ))
    }

    /// Minimum expression given fixed miscalls.
    pub fn min_expression(&self, miscalls_exact: &Miscalls, miscalls_mismatch: &Miscalls) -> u32 {
        self.counts.exact.saturating_sub(miscalls_exact.total_to(self.feature_id)) +
        self.counts.mismatch.saturating_sub(miscalls_mismatch.total_to(self.feature_id)) +
        miscalls_exact.total_from(self.feature_id) + miscalls_mismatch.total_from(self.feature_id)
    }
}

impl AbstractFeatureModel for FeatureModel {
    fn feature_id(&self) -> FeatureID {
        self.feature_id
    }

    fn neighbors(&self) -> &[FeatureID] {
        &self.neighbors
    }

    fn min_dist(&self) -> u8 {
        self.min_dist
    }

    fn event_count_non_dropout(
        &self,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> u32 {
        let calls_exact = self.counts.exact.saturating_sub(
            miscalls_exact.total_to(self.feature_id)
        );
        let calls_mismatch = self.counts.mismatch.saturating_sub(
            miscalls_mismatch.total_to(self.feature_id)
        );

        calls_exact + calls_mismatch +
        miscalls_exact.total_from(self.feature_id) +
        miscalls_mismatch.total_from(self.feature_id)
    }

    fn prob_dropout(&self) -> f64 {
        self.event_probs[2]
    }

    fn prob_miscall_exact(&self, neighbor_index: usize) -> f64 {
        self.event_probs[3 + neighbor_index]
    }

    fn prob_miscall_mismatch(&self, neighbor_index: usize) -> f64 {
        self.event_probs[3 + self.neighbors.len() + neighbor_index]
    }
}

pub struct NoiseModel {
    pub(crate) feature_id: FeatureID,
    pub(crate) not_expressed_feature_ids: Vec<FeatureID>,
    pub(crate) not_expressed_counts: Vec<Counts>,
    probs_miscall_exact: Vec<Prob>,
    probs_miscall_mismatch: Vec<Prob>,
    prob_dropout: Prob,
    neighbors: Vec<FeatureID>,
    min_dist: u8,
}

impl NoiseModel {
    pub fn new(
        feature_id: FeatureID,
        not_expressed_feature_ids: Vec<FeatureID>,
        not_expressed_counts: Vec<Counts>,
        codebook: &Codebook,
        xi: &Xi
    ) -> Self {
        assert!(
            codebook.min_dist == 2 || codebook.min_dist == 4,
            "unsupported hamming distance: {}, supported: 2, 4",
            codebook.min_dist
        );

        // miscalls generated from noise
        let xi_miscall = codebook.records().iter().filter_map(|feat_rec| {
            if feat_rec.expressed() {
                Some(xi.prob(codebook.noise().codeword(), feat_rec.codeword()))
            } else { None }
        }).collect_vec();
        let neighbors = codebook.records().iter().filter_map(|rec| {
            if rec.expressed() {
                Some(codebook.get_id(rec.name()))
            } else { None }
        }).collect_vec();

        let probs_miscall_exact = xi_miscall.iter().map(|xi| xi[0]).collect_vec();
        let probs_miscall_mismatch = if codebook.min_dist == 4 {
            xi_miscall.iter().map(|xi| xi[1]).collect_vec()
        } else {
            vec![]
        };

        let prob_total_miscall_exact = LogProb::ln_sum_exp(&probs_miscall_exact);
        let prob_total_miscall_mismatch = LogProb::ln_sum_exp(&probs_miscall_mismatch);

        // probs for calls of not expressed features
        let xi_not_expressed = not_expressed_feature_ids.iter().map(|&f| {
            xi.prob(codebook.noise().codeword(), codebook.record(f).codeword())
        }).collect_vec();

        let prob_not_expressed_exact = LogProb::ln_sum_exp(
            &xi_not_expressed.iter().map(|xi| xi[0]).collect_vec()
        );
        let prob_not_expressed_mismatch = LogProb::ln_sum_exp(
            &xi_not_expressed.iter().map(|xi| xi[1]).collect_vec()
        );

        // dropout
        let prob_dropout = LogProb::ln_sum_exp(&[
            prob_not_expressed_exact,
            prob_not_expressed_mismatch,
            prob_total_miscall_exact,
            prob_total_miscall_mismatch
        ]).ln_one_minus_exp();

        NoiseModel {
            feature_id: feature_id,
            not_expressed_feature_ids: not_expressed_feature_ids,
            not_expressed_counts: not_expressed_counts,
            probs_miscall_exact: probs_miscall_exact.iter().map(|&p| Prob::from(p)).collect_vec(),
            probs_miscall_mismatch: probs_miscall_mismatch.iter().map(|&p| Prob::from(p)).collect_vec(),
            prob_dropout: Prob::from(prob_dropout),
            neighbors: neighbors,
            min_dist: codebook.min_dist
        }
    }
}

impl AbstractFeatureModel for NoiseModel {
    fn feature_id(&self) -> FeatureID {
        self.feature_id
    }

    fn neighbors(&self) -> &[FeatureID] {
        &self.neighbors
    }

    fn min_dist(&self) -> u8 {
        self.min_dist
    }

    fn event_count_non_dropout(
        &self,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> u32 {
        let mut calls_exact = 0;
        let mut calls_mismatch = 0;
        for (&id, counts) in self.not_expressed_feature_ids.iter().zip(&self.not_expressed_counts) {
            calls_exact += counts.exact.saturating_sub(miscalls_exact.total_to(id));
            calls_mismatch += counts.mismatch.saturating_sub(miscalls_mismatch.total_to(id));
        }

        let event_count = calls_exact + calls_mismatch +
                          miscalls_exact.total_from(self.feature_id) +
                          miscalls_mismatch.total_from(self.feature_id);

        event_count
    }

    fn prob_dropout(&self) -> f64 {
        //*self.prob_dropout
        1.0
    }

    fn prob_miscall_exact(&self, neighbor_index: usize) -> f64 {
        *self.probs_miscall_exact[neighbor_index]
    }

    fn prob_miscall_mismatch(&self, neighbor_index: usize) -> f64 {
        *self.probs_miscall_mismatch[neighbor_index]
    }
}
<|endoftext|>