<commit_before>extern crate ndarray;

use ndarray_ext::NdArray;
use op;
use std::collections::BTreeMap;
use std::collections::btree_map::Entry;
use std::mem;
use std::rc::Rc;
use tensor::Tensor;

#[derive(Clone, Debug, PartialEq)]
pub enum EvaluationError {
    ComputeFailed {
        node_name: String,
        msg: String
    },
    NoOutput
}

// Context in evaluation of `node`
pub struct OpComputeContext<'a, 'b>
{
    pub node: &'a Tensor, // Expose to its op for convenience
    xs: Vec<&'b NdArray>,
}

impl<'a, 'b> OpComputeContext<'a, 'b>
{
    #[inline]
    pub fn grab_inputs(&self) -> &[&NdArray]
    {
        self.xs.as_slice()
    }

    #[inline]
    #[allow(mutable_transmutes)]
    pub unsafe fn grab_assignable_inputs(&mut self) -> &mut [&mut NdArray]
    {
        mem::transmute(self.xs.as_slice())
    }
}

#[inline]
// Grabs input arrays for `node`'s evaluation.
// Returns "None" when failed to aggregate even one of input arrays.
fn _grab_input_arrays<'a, 'b: 'a>(node: &'b Tensor,
                                  store: &'a ResourceStore,
                                  feed_store: &FeedStore<'a>) -> Option<Vec<&'a NdArray>> {
    #[inline]
    fn find_array_of<'a, 'b: 'a>(
        x: &'b Tensor,
        store: &'a ResourceStore,
        feed_store: &FeedStore<'a>,
        value_index: usize,
    ) -> Option<&'a NdArray>
    {
        if let Some(ref per) = x.persistent_array {
            Some(per.get_array())
        } else if x.is_placeholder {
            Some(feed_store[x.resource_lookup_key.get()])
        } else {
            match store[x.resource_lookup_key.get()].value[value_index] {
                Ok(ref a) => Some(a),
                // hoping for x.inputs[i] to have the value
                Err(::op::ComputeError::Delegate { to: i }) => {
                    find_array_of(&x.inputs[i], store, feed_store, x.input_indices[i])
                },
                _ => None,  // None for hopeless errors
            }
        }
    }

    let input_nodes = &node.inputs;
    let mut input_arrays = Vec::with_capacity(input_nodes.len());
    for (x, &i) in input_nodes.into_iter().zip(node.input_indices.iter()) {
        if let Some(res) = find_array_of(x, store, feed_store, i) {
            input_arrays.push(res);
        } else {
            // Early return
            return None
        }
    }
    Some(input_arrays)
}

struct NodeWithValue<'a>
{
    node:  &'a Tensor,
    value: op::ComputeResult,
    // How many resources of this node does user requires.
    // When this is reduced to one, `value` is ready to be moved out (without copy).
    pending_count: usize,
}

impl<'a> Tensor
{
    #[inline]
    fn with_value(&'a self, val: op::ComputeResult) -> NodeWithValue<'a>
    {
        NodeWithValue {
            node:          self,
            value:         val,
            pending_count: 0,
        }
    }
}

/// Evaluates given symbolic tensors.
///
/// ```
/// extern crate ndarray;
/// extern crate autograd as ag;
///
/// let ref a = ag::zeros(&[2]);
/// let ref b = ag::ones(&[2]);
///
/// // eval two tensors at once.
/// let evaluated = ag::eval(&[a, b], &[]);
/// assert_eq!(evaluated[0], Ok(ndarray::arr1(&[0., 0.]).into_dyn()));
/// assert_eq!(evaluated[1], Ok(ndarray::arr1(&[1., 1.]).into_dyn()));
/// ```
// FIXME: Annoying lifetime params
pub fn eval<'a, 'b: 'a, 'c: 'a, T, U>(
    tensors: &[T],
    feeds: U,
) -> Vec<Result<ndarray::Array<f32, ndarray::IxDyn>, EvaluationError>>
where
    T: AsRef<Tensor>,
    U: IntoIterator<Item = &'a (&'b Tensor, &'c ndarray::Array<f32, ndarray::IxDyn>)>,
{
    // Run graph
    let feeds = feeds.into_iter().collect::<Vec<_>>();
    let mut output_storage = eval_internal(&tensors.iter().map(|t| t.as_ref()).collect(), &feeds);

    // Treat in-place or delegation ops
    let creators = tensors
        .iter()
        .map(|x| {
            let x = x.as_ref();
            let creator = if x.is_placeholder || x.persistent_array.is_some() {
                x
            } else {
                let creator = find_resource_creator(&output_storage, x);
                output_storage[creator.resource_lookup_key.get()].pending_count += 1;
                creator
            };
            creator
        })
        .collect::<Vec<&Tensor>>();

    // Shrink to fit (output_storage is moved)
    let mut key2res: BTreeMap<usize, NodeWithValue> = finalize_resource_store(output_storage);

    // Aggregate return values
    creators
        .iter()
        .map(|ref creator| {
            if let Some(ref per) = creator.persistent_array {
                // Rarely happens (case that a persistent array given by user is required)
                Ok(per.get_array().clone())
            } else if creator.is_placeholder {
                // Rarely happens (case that a feed by user is required)
                Ok(get_fed_resource(creator, &feeds).clone())
            } else {
                let res = match key2res.entry(creator.resource_lookup_key.get()) {
                    Entry::Occupied(mut ent) => {
                        if ent.get().pending_count == 1 { // move out the resource.
                            let mut got = ent.remove();
                            // Using first item of the compute result
                            map_err(got.value.remove(0), creator)
                        } else { // "clone" the resource.
                            let mut got = ent.get_mut();
                            got.pending_count -= 1;
                            // Using first item of the compute result
                            map_err(got.value[0].clone(), creator)
                        }
                    }
                    _ => unreachable!(),
                };
                res
            }
        })
        .collect()
}

/// Runs given symbolic tensors.
///
/// Runs, but returns nothing.
///
/// ```
/// extern crate ndarray;
/// extern crate autograd as ag;
///
/// let a = ag::variable(ndarray::arr1(&[1., 1.]));
/// let b = ag::ones(&[2]);
/// let c = ag::sub_inplace(a, &b);
///
/// // runs inplace op.
/// ag::run(&[&c], &[]);
/// ```
pub fn run<'a, 'b: 'a, 'c: 'a, T, U>(tensors: &[T], feeds: U)
where
    T: AsRef<Tensor>,
    U: IntoIterator<Item = &'a (&'b Tensor, &'c NdArray)>,
{
    // Just run the graph
    eval_internal(
        &tensors.iter().map(|t| t.as_ref()).collect(),
        &feeds.into_iter().collect(),
    );
}

#[inline]
// Recursive function which seeks a node holding the x's resource.
// Actual recursion "rarely" happens.
fn find_resource_creator<'a, 'b>(storage: &ResourceStore, x: &'b Tensor) -> &'b Tensor
{
    match storage[x.resource_lookup_key.get()].value[0] {
        Err(::op::ComputeError::Delegate { to: i }) => {
            find_resource_creator(storage, &x.inputs[i])
        }
        _ => x
    }
}

type EvalResult = Result<NdArray, EvaluationError>;

#[inline]
fn map_err<'a>(res: Result<NdArray, ::op::ComputeError>,
               node: &'a Tensor) -> EvalResult {
    match res {
        Ok(arr) => Ok(arr),
        Err(::op::ComputeError::BadInput(ref msg)) =>
            Err(EvaluationError::ComputeFailed {
                node_name: node.op.name().to_string(),
                msg: msg.to_string() }
            ),
        Err(::op::ComputeError::NoOutput) => Err(EvaluationError::NoOutput),
        _ => unreachable!()
    }
}

type ResourceStore<'a> = Vec<NodeWithValue<'a>>;
type FeedStore<'a> = Vec<&'a NdArray>;

#[inline]
// TODO: Use raw pointer comparison after removing "Rc"
fn get_fed_resource<'a>(node: &Tensor, feeds: &Vec<&(&Tensor, &'a NdArray)>) -> &'a NdArray
{
    // Linear search is suitable because number of feeds are so small in most cases
    for feed in feeds {
        if Rc::ptr_eq(feed.0, node) {
            return feed.1;
        }
    }
    // panic
    panic!("Placeholder unfilled. See backtrace.");
}

// Actually evaluates nodes "targets".
fn eval_internal<'a>(
    targets: &Vec<&'a Tensor>,
    feeds: &Vec<&(&'a Tensor, &NdArray)>,
) -> ResourceStore<'a>
{
    let mut resource_store = Vec::new();
    let mut feed_store = Vec::new();

    // Obtain array resources while visiting nodes in topological order.
    // Stack-based DFS is used, to avoid stack overflow in explicit recursion.
    let mut dfs_stack: Vec<(&Tensor, bool)> = targets.iter().map(|&x| (x, false)).collect();
    while let Some((node, is_parent)) = dfs_stack.pop() {
        if is_parent {
            // Visit this node
            if node.is_placeholder {
                node.resource_lookup_key.set(feed_store.len());
                feed_store.push(get_fed_resource(node, &feeds));
            } else {
                node.resource_lookup_key.set(resource_store.len());
                if node.persistent_array.is_none() {
                    let y = if let Some(xs) = _grab_input_arrays(node, &resource_store, &feed_store) {
                        node.op.compute(OpComputeContext { node, xs })
                    } else {
                        vec![Err(::op::ComputeError::Delegate {to: 0})]
                    };
                    resource_store.push(node.with_value(y));
                }
            }
        } else {
            // Update dfs stack
            dfs_stack.push((node, true));
            // Push children if needed
            for child in &node.inputs {
                // TODO: Use raw pointer comparison after removing "Rc"
                let visited = {
                    let k = child.resource_lookup_key.get();
                    k < resource_store.len() && Rc::ptr_eq(child, resource_store[k].node)
                };
                if !visited {
                    dfs_stack.push((child, false));
                }
            }
        }
    }
    resource_store
}

// Shrink it by dropping useless resources
// and convert it into mappings of {lookup key => resource}
fn finalize_resource_store(mut vec: ResourceStore) -> BTreeMap<usize, NodeWithValue>
{
    let mut retained_keys = Vec::new();
    let len = vec.len();
    {
        let mut_ref = &mut vec;
        let mut del = 0;
        // Align the resources. Unused resources are placed in the latter part.
        {
            let v = &mut **mut_ref;

            for i in 0..len {
                if v[i].pending_count == 0 {
                    del += 1;
                    continue;
                }
                retained_keys.push(i);
                if del > 0 {
                    // slides i th node forward
                    v.swap(i - del, i);
                }
            }
        }
        if del > 0 {
            // Drop unused resources here
            mut_ref.truncate(len - del);
        }
    }
    debug_assert_eq!(vec.len(), retained_keys.len());
    // `retained_keys` are sorted.
    // `vec` is moved into Map.
    retained_keys.into_iter().zip(vec).collect()
}

#[test]
fn test_eval()
{
    let ref v = ::ops::placeholder(&[3, 2, 1]);
    let ref z = ::ops::reduce_sum(&::ops::squeeze(v, &[2]), &[0, 1], false);
    let ref g = ::ops::grad(&[z], &[v]);
    let eval_result = &eval(g, &[(v, &::ndarray_ext::ones(&[3, 2, 1]))])[0];
    assert_eq!(eval_result.as_ref().unwrap().shape(), &[3, 2, 1]);
}

#[test]
fn test_constant_eval()
{
    let arr = ndarray::arr1(&[0., 0., 0.]);
    assert_eq!(Ok(arr.clone().into_dyn()), ::variable(arr).eval(&[]));
}

#[test]
fn test_placeholder_eval()
{
    let arr = ::ndarray_ext::ones(&[3, 2, 1]);
    let ref v = ::ops::placeholder(&[3, 2, 1]);
    let eval_result = eval(&[v], &[(v, &arr)]);
    assert_eq!(eval_result[0], Ok(arr));
}

#[test]
fn test_eval_internal()
{
    let ref v = ::ops::placeholder(&[3, 2, 1]);
    let ref z = ::ops::squeeze(v, &[2]);
    let ref g = ::ops::grad_with_default(&[z], &[v], &[&::ones(&z.shape())]);
    let storage = eval_internal(&vec![&g[0]], &vec![&(v, &::ndarray_ext::ones(&[3, 2, 1]))]);

    assert_eq!(
        storage.iter().map(|x| x.node.op.name()).collect::<Vec<_>>(),
        vec![
            "ConvertToTensor",
            "Squeeze", // forward end
            "Shape",
            "Ones",
            "ExpandDims",
        ]
    );
}
<commit_msg>Runtime errors in ag::run cause panics just before exit<commit_after>extern crate ndarray;

use ndarray_ext::NdArray;
use op;
use std::collections::BTreeMap;
use std::collections::btree_map::Entry;
use std::mem;
use std::rc::Rc;
use tensor::Tensor;

#[derive(Clone, Debug, PartialEq)]
pub enum EvaluationError {
    ComputeFailed {
        node_name: String,
        msg: String
    },
    NoOutput
}

// Context in evaluation of `node`
pub struct OpComputeContext<'a, 'b>
{
    pub node: &'a Tensor, // Expose to its op for convenience
    xs: Vec<&'b NdArray>,
}

impl<'a, 'b> OpComputeContext<'a, 'b>
{
    #[inline]
    pub fn grab_inputs(&self) -> &[&NdArray]
    {
        self.xs.as_slice()
    }

    #[inline]
    #[allow(mutable_transmutes)]
    pub unsafe fn grab_assignable_inputs(&mut self) -> &mut [&mut NdArray]
    {
        mem::transmute(self.xs.as_slice())
    }
}

#[inline]
// Grabs input arrays for `node`'s evaluation.
// Returns "None" when failed to aggregate even one of input arrays.
fn _grab_input_arrays<'a, 'b: 'a>(node: &'b Tensor,
                                  store: &'a ResourceStore,
                                  feed_store: &FeedStore<'a>) -> Option<Vec<&'a NdArray>> {
    #[inline]
    fn find_array_of<'a, 'b: 'a>(
        x: &'b Tensor,
        store: &'a ResourceStore,
        feed_store: &FeedStore<'a>,
        value_index: usize,
    ) -> Option<&'a NdArray>
    {
        if let Some(ref per) = x.persistent_array {
            Some(per.get_array())
        } else if x.is_placeholder {
            Some(feed_store[x.resource_lookup_key.get()])
        } else {
            match store[x.resource_lookup_key.get()].value[value_index] {
                Ok(ref a) => Some(a),
                // hoping for x.inputs[i] to have the value
                Err(::op::ComputeError::Delegate { to: i }) => {
                    find_array_of(&x.inputs[i], store, feed_store, x.input_indices[i])
                },
                _ => None,  // None for hopeless errors
            }
        }
    }

    let input_nodes = &node.inputs;
    let mut input_arrays = Vec::with_capacity(input_nodes.len());
    for (x, &i) in input_nodes.into_iter().zip(node.input_indices.iter()) {
        if let Some(res) = find_array_of(x, store, feed_store, i) {
            input_arrays.push(res);
        } else {
            // Early return
            return None
        }
    }
    Some(input_arrays)
}

struct NodeWithValue<'a>
{
    node:  &'a Tensor,
    value: op::ComputeResult,
    // How many resources of this node does user requires.
    // When this is reduced to one, `value` is ready to be moved out (without copy).
    pending_count: usize,
}

impl<'a> Tensor
{
    #[inline]
    fn with_value(&'a self, val: op::ComputeResult) -> NodeWithValue<'a>
    {
        NodeWithValue {
            node:          self,
            value:         val,
            pending_count: 0,
        }
    }
}

/// Evaluates given symbolic tensors.
///
/// ```
/// extern crate ndarray;
/// extern crate autograd as ag;
///
/// let ref a = ag::zeros(&[2]);
/// let ref b = ag::ones(&[2]);
///
/// // eval two tensors at once.
/// let evaluated = ag::eval(&[a, b], &[]);
/// assert_eq!(evaluated[0], Ok(ndarray::arr1(&[0., 0.]).into_dyn()));
/// assert_eq!(evaluated[1], Ok(ndarray::arr1(&[1., 1.]).into_dyn()));
/// ```
// FIXME: Annoying lifetime params
pub fn eval<'a, 'b: 'a, 'c: 'a, T, U>(
    tensors: &[T],
    feeds: U,
) -> Vec<Result<ndarray::Array<f32, ndarray::IxDyn>, EvaluationError>>
where
    T: AsRef<Tensor>,
    U: IntoIterator<Item = &'a (&'b Tensor, &'c ndarray::Array<f32, ndarray::IxDyn>)>,
{
    // Run graph
    let feeds = feeds.into_iter().collect::<Vec<_>>();
    let mut output_storage = eval_internal(&tensors.iter().map(|t| t.as_ref()).collect(), &feeds);

    // Treat in-place or delegation ops
    let creators = tensors
        .iter()
        .map(|x| {
            let x = x.as_ref();
            let creator = if x.is_placeholder || x.persistent_array.is_some() {
                x
            } else {
                let creator = find_resource_creator(&output_storage, x);
                output_storage[creator.resource_lookup_key.get()].pending_count += 1;
                creator
            };
            creator
        })
        .collect::<Vec<&Tensor>>();

    // Shrink to fit (output_storage is moved)
    let mut key2res: BTreeMap<usize, NodeWithValue> = finalize_resource_store(output_storage);

    // Aggregate return values
    creators
        .iter()
        .map(|ref creator| {
            if let Some(ref per) = creator.persistent_array {
                // Rarely happens (case that a persistent array given by user is required)
                Ok(per.get_array().clone())
            } else if creator.is_placeholder {
                // Rarely happens (case that a feed by user is required)
                Ok(get_fed_resource(creator, &feeds).clone())
            } else {
                let res = match key2res.entry(creator.resource_lookup_key.get()) {
                    Entry::Occupied(mut ent) => {
                        if ent.get().pending_count == 1 { // move out the resource.
                            let mut got = ent.remove();
                            // Using first item of the compute result
                            map_err(got.value.remove(0), creator)
                        } else { // "clone" the resource.
                            let mut got = ent.get_mut();
                            got.pending_count -= 1;
                            // Using first item of the compute result
                            map_err(got.value[0].clone(), creator)
                        }
                    }
                    _ => unreachable!(),
                };
                res
            }
        })
        .collect()
}

/// Runs given symbolic tensors.
///
/// Runs, but returns nothing.
///
/// ```
/// extern crate ndarray;
/// extern crate autograd as ag;
///
/// let a = ag::variable(ndarray::arr1(&[1., 1.]));
/// let b = ag::ones(&[2]);
/// let c = ag::sub_inplace(a, &b);
///
/// // runs inplace op.
/// ag::run(&[&c], &[]);
/// ```
pub fn run<'a, 'b: 'a, 'c: 'a, T, U>(tensors: &[T], feeds: U)
where
    T: AsRef<Tensor>,
    U: IntoIterator<Item = &'a (&'b Tensor, &'c NdArray)>,
{
    // Just run the graph
    let ret = eval_internal(
        &tensors.iter().map(|t| t.as_ref()).collect(),
        &feeds.into_iter().collect(),
    );
    // Panic as necessary
    for r in ret {
        panic_if_user_err(&r.value[0])
    }
}

#[inline]
fn panic_if_user_err(res: &Result<NdArray, op::ComputeError>)
{
    match res {
        &Err(op::ComputeError::BadInput(_)) => {
            res.as_ref().unwrap();
        },
        _ => {},
    }
}

#[inline]
// Recursive function which seeks a node holding the x's resource.
// Actual recursion "rarely" happens.
fn find_resource_creator<'a, 'b>(storage: &ResourceStore, x: &'b Tensor) -> &'b Tensor
{
    match storage[x.resource_lookup_key.get()].value[0] {
        Err(::op::ComputeError::Delegate { to: i }) => {
            find_resource_creator(storage, &x.inputs[i])
        }
        _ => x
    }
}

type EvalResult = Result<NdArray, EvaluationError>;

#[inline]
fn map_err<'a>(res: Result<NdArray, ::op::ComputeError>,
               node: &'a Tensor) -> EvalResult {
    match res {
        Ok(arr) => Ok(arr),
        Err(::op::ComputeError::BadInput(ref msg)) =>
            Err(EvaluationError::ComputeFailed {
                node_name: node.op.name().to_string(),
                msg: msg.to_string() }
            ),
        Err(::op::ComputeError::NoOutput) => Err(EvaluationError::NoOutput),
        _ => unreachable!()
    }
}

type ResourceStore<'a> = Vec<NodeWithValue<'a>>;
type FeedStore<'a> = Vec<&'a NdArray>;

#[inline]
// TODO: Use raw pointer comparison after removing "Rc"
fn get_fed_resource<'a>(node: &Tensor, feeds: &Vec<&(&Tensor, &'a NdArray)>) -> &'a NdArray
{
    // Linear search is suitable because number of feeds are so small in most cases
    for feed in feeds {
        if Rc::ptr_eq(feed.0, node) {
            return feed.1;
        }
    }
    // panic
    panic!("Placeholder unfilled. See backtrace.");
}

// Actually evaluates nodes "targets".
fn eval_internal<'a>(
    targets: &Vec<&'a Tensor>,
    feeds: &Vec<&(&'a Tensor, &NdArray)>,
) -> ResourceStore<'a>
{
    let mut resource_store = Vec::new();
    let mut feed_store = Vec::new();

    // Obtain array resources while visiting nodes in topological order.
    // Stack-based DFS is used, to avoid stack overflow in explicit recursion.
    let mut dfs_stack: Vec<(&Tensor, bool)> = targets.iter().map(|&x| (x, false)).collect();
    while let Some((node, is_parent)) = dfs_stack.pop() {
        if is_parent {
            // Visit this node
            if node.is_placeholder {
                node.resource_lookup_key.set(feed_store.len());
                feed_store.push(get_fed_resource(node, &feeds));
            } else {
                node.resource_lookup_key.set(resource_store.len());
                if node.persistent_array.is_none() {
                    let y = if let Some(xs) = _grab_input_arrays(node, &resource_store, &feed_store) {
                        node.op.compute(OpComputeContext { node, xs })
                    } else {
                        vec![Err(::op::ComputeError::Delegate {to: 0})]
                    };
                    resource_store.push(node.with_value(y));
                }
            }
        } else {
            // Update dfs stack
            dfs_stack.push((node, true));
            // Push children if needed
            for child in &node.inputs {
                // TODO: Use raw pointer comparison after removing "Rc"
                let visited = {
                    let k = child.resource_lookup_key.get();
                    k < resource_store.len() && Rc::ptr_eq(child, resource_store[k].node)
                };
                if !visited {
                    dfs_stack.push((child, false));
                }
            }
        }
    }
    resource_store
}

// Shrink it by dropping useless resources
// and convert it into mappings of {lookup key => resource}
fn finalize_resource_store(mut vec: ResourceStore) -> BTreeMap<usize, NodeWithValue>
{
    let mut retained_keys = Vec::new();
    let len = vec.len();
    {
        let mut_ref = &mut vec;
        let mut del = 0;
        // Align the resources. Unused resources are placed in the latter part.
        {
            let v = &mut **mut_ref;

            for i in 0..len {
                if v[i].pending_count == 0 {
                    del += 1;
                    continue;
                }
                retained_keys.push(i);
                if del > 0 {
                    // slides i th node forward
                    v.swap(i - del, i);
                }
            }
        }
        if del > 0 {
            // Drop unused resources here
            mut_ref.truncate(len - del);
        }
    }
    debug_assert_eq!(vec.len(), retained_keys.len());
    // `retained_keys` are sorted.
    // `vec` is moved into Map.
    retained_keys.into_iter().zip(vec).collect()
}

#[test]
fn test_eval()
{
    let ref v = ::ops::placeholder(&[3, 2, 1]);
    let ref z = ::ops::reduce_sum(&::ops::squeeze(v, &[2]), &[0, 1], false);
    let ref g = ::ops::grad(&[z], &[v]);
    let eval_result = &eval(g, &[(v, &::ndarray_ext::ones(&[3, 2, 1]))])[0];
    assert_eq!(eval_result.as_ref().unwrap().shape(), &[3, 2, 1]);
}

#[test]
fn test_constant_eval()
{
    let arr = ndarray::arr1(&[0., 0., 0.]);
    assert_eq!(Ok(arr.clone().into_dyn()), ::variable(arr).eval(&[]));
}

#[test]
fn test_placeholder_eval()
{
    let arr = ::ndarray_ext::ones(&[3, 2, 1]);
    let ref v = ::ops::placeholder(&[3, 2, 1]);
    let eval_result = eval(&[v], &[(v, &arr)]);
    assert_eq!(eval_result[0], Ok(arr));
}

#[test]
fn test_eval_internal()
{
    let ref v = ::ops::placeholder(&[3, 2, 1]);
    let ref z = ::ops::squeeze(v, &[2]);
    let ref g = ::ops::grad_with_default(&[z], &[v], &[&::ones(&z.shape())]);
    let storage = eval_internal(&vec![&g[0]], &vec![&(v, &::ndarray_ext::ones(&[3, 2, 1]))]);

    assert_eq!(
        storage.iter().map(|x| x.node.op.name()).collect::<Vec<_>>(),
        vec![
            "ConvertToTensor",
            "Squeeze", // forward end
            "Shape",
            "Ones",
            "ExpandDims",
        ]
    );
}
<|endoftext|>


<commit_before>use core::prelude::*;

use core::atomic::{AtomicUsize, Ordering, ATOMIC_USIZE_INIT};
use core::cell::UnsafeCell;
use core::ops::{Deref, DerefMut};

/// A reader-writer lock
///
/// This type of lock allows a number of readers or at most one writer at any
/// point in time. The write portion of this lock typically allows modification
/// of the underlying data (exclusive access) and the read portion of this lock
/// typically allows for read-only access (shared access).
///
/// The type parameter `T` represents the data that this lock protects. It is
/// required that `T` satisfies `Send` to be shared across tasks and `Sync` to
/// allow concurrent access through readers. The RAII guards returned from the
/// locking methods implement `Deref` (and `DerefMut` for the `write` methods)
/// to allow access to the contained of the lock.
///
/// Based on
/// https://jfdube.wordpress.com/2014/01/03/implementing-a-recursive-read-write-spinlock/
///
pub struct RwLock<T>
{
    lock: AtomicUsize,
    data: UnsafeCell<T>,
}

/// A guard to which the protected data can be read
///
/// When the guard falls out of scope it will decrement the read count,
/// potentially releasing the lock.
pub struct RwLockReadGuard<'a, T:'a>
{
    lock: &'a AtomicUsize,
    data: &'a T,
}

/// A guard to which the protected data can be written
///
/// When the guard falls out of scope it will release the lock.
pub struct RwLockWriteGuard<'a, T:'a>
{
    lock: &'a AtomicUsize,
    data: &'a mut T,
}

unsafe impl<T> Sync for RwLock<T> {}
unsafe impl<T:'static+Send> Send for RwLock<T> {}

/// A RwLock which may be used statically.
///
/// ```
/// use spin::{StaticRwLock, STATIC_RWLOCK_INIT};
///
/// static SPRWLCK: StaticRwLock = STATIC_RWLOCK_INIT;
///
/// fn demo() {
///     let lock = SPRWLCK.read();
///     // do something with lock
///     drop(lock);
/// }
/// ```
pub type StaticRwLock = RwLock<()>;

/// A initializer for StaticRwLock, containing no data.
pub const STATIC_RWLOCK_INIT: StaticRwLock = RwLock {
    lock: ATOMIC_USIZE_INIT,
    data: UnsafeCell { value: () },
};

const USIZE_MSB: usize = ::core::isize::MIN as usize;

impl<T> RwLock<T>
{
    /// Creates a new spinlock wrapping the supplied data.
    #[inline]
    pub fn new(user_data: T) -> RwLock<T>
    {
        RwLock
        {
            lock: ATOMIC_USIZE_INIT,
            data: UnsafeCell::new(user_data),
        }
    }

    /// Locks this rwlock with shared read access, blocking the current thread
    /// until it can be acquired.
    ///
    /// The calling thread will be blocked until there are no more writers which
    /// hold the lock. There may be other readers currently inside the lock when
    /// this method returns. This method does not provide any guarantees with
    /// respect to the ordering of whether contentious readers or writers will
    /// acquire the lock first.
    ///
    /// Returns an RAII guard which will release this thread's shared access
    /// once it is dropped.
    ///
    /// ```
    /// let mylock = spin::RwLock::new(0us);
    /// {
    ///     let mut data = mylock.read();
    ///     // The lock is now locked and the data can be read
    ///     println!("{}", *data);
    ///     // The lock is dropped
    /// }
    /// ```
    #[inline]
    pub fn read<'a>(&'a self) -> RwLockReadGuard<'a, T>
    {
        // (funny do-while loop)
        while {
            // Old value, with write bit unset
            let mut old;

            // Wait for for writer to go away before doing expensive atomic ops
            // (funny do-while loop)
            while {
                old = self.lock.load(Ordering::Relaxed);
                old & USIZE_MSB != 0
            } {}

            // unset write bit
            old &= !USIZE_MSB;

            let new = old + 1;
            debug_assert!(new != (!USIZE_MSB) & (-1));

            self.lock.compare_and_swap(old, new, Ordering::SeqCst) != old
        } {}
        RwLockReadGuard {
            lock: &self.lock,
            data: unsafe { & *self.data.get() },
        }
    }

    /// Attempt to acquire this lock with shared read access.
    ///
    /// This function will never block and will return immediately if `read`
    /// would otherwise succeed. Returns `Some` of an RAII guard which will
    /// release the shared access of this thread when dropped, or `None` if the
    /// access could not be granted. This method does not provide any
    /// guarantees with respect to the ordering of whether contentious readers
    /// or writers will acquire the lock first.
    ///
    /// ```
    /// let mylock = spin::RwLock::new(0us);
    /// {
    ///     if let Some(data) = mylock.try_read() {
    ///         // The lock is now locked and the data can be read
    ///         println!("{}", *data);
    ///         // The lock is dropped
    ///     };
    /// }
    /// ```
    #[inline]
    pub fn try_read(&self) -> Option<RwLockReadGuard<T>>
    {
        // Old value, with write bit unset
        let old = (!USIZE_MSB) & self.lock.load(Ordering::Relaxed);

        let new = old + 1;
        debug_assert!(new != (!USIZE_MSB) & (-1));
        if self.lock.compare_and_swap(old,
                                      new,
                                      Ordering::SeqCst) == old
        {
            Some(RwLockReadGuard {
                lock: &self.lock,
                data: unsafe { & *self.data.get() },
            })
        } else {
            None
        }
    }

    /// Lock this rwlock with exclusive write access, blocking the current
    /// thread until it can be acquired.
    ///
    /// This function will not return while other writers or other readers
    /// currently have access to the lock.
    ///
    /// Returns an RAII guard which will drop the write access of this rwlock
    /// when dropped.
    ///
    /// ```
    /// let mylock = spin::RwLock::new(0us);
    /// {
    ///     let mut data = mylock.write();
    ///     // The lock is now locked and the data can be written
    ///     *data += 1;
    ///     // The lock is dropped
    /// }
    /// ```
    #[inline]
    pub fn write<'a>(&'a self) -> RwLockWriteGuard<'a, T>
    {
        loop
        {
            // Old value, with write bit unset.
            let old = (!USIZE_MSB) & self.lock.load(Ordering::Relaxed);
            // Old value, with write bit set.
            let new = USIZE_MSB | old;
            if self.lock.compare_and_swap(old,
                                          new,
                                          Ordering::SeqCst) == old
            {
                // Wait for readers to go away, then lock is ours.
                while self.lock.load(Ordering::Relaxed) != USIZE_MSB { }
                break
            }
        }
        RwLockWriteGuard {
            lock: &self.lock,
            data: unsafe { &mut *self.data.get() },
        }
    }

    /// Attempt to lock this rwlock with exclusive write access.
    ///
    /// This function does not ever block, and it will return `None` if a call
    /// to `write` would otherwise block. If successful, an RAII guard is
    /// returned.
    ///
    /// ```
    /// let mylock = spin::RwLock::new(0us);
    /// {
    ///     if let Some(mut data) = mylock.try_write() {
    ///         // The lock is now locked and the data can be written
    ///         *data += 1;
    ///         // The lock is implicitly dropped
    ///     };
    /// }
    /// ```
    #[inline]
    pub fn try_write(&self) -> Option<RwLockWriteGuard<T>>
    {
        if self.lock.compare_and_swap(0,
                                      USIZE_MSB,
                                      Ordering::SeqCst) == 0
        {
            Some(RwLockWriteGuard {
                lock: &self.lock,
                data: unsafe { &mut *self.data.get() },
            })
        } else {
            None
        }
    }
}

impl<'rwlock, T> Deref for RwLockReadGuard<'rwlock, T> {
    type Target = T;

    fn deref(&self) -> &T { self.data }
}

impl<'rwlock, T> Deref for RwLockWriteGuard<'rwlock, T> {
    type Target = T;

    fn deref(&self) -> &T { self.data }
}

impl<'rwlock, T> DerefMut for RwLockWriteGuard<'rwlock, T> {
    fn deref_mut(&mut self) -> &mut T { self.data }
}

#[unsafe_destructor]
impl<'rwlock, T> Drop for RwLockReadGuard<'rwlock, T> {
    fn drop(&mut self) {
        debug_assert!(self.lock.load(Ordering::Relaxed) & (!USIZE_MSB) > 0);
        self.lock.fetch_sub(1, Ordering::SeqCst);
    }
}

#[unsafe_destructor]
impl<'rwlock, T> Drop for RwLockWriteGuard<'rwlock, T> {
    fn drop(&mut self) {
        debug_assert_eq!(self.lock.load(Ordering::Relaxed), USIZE_MSB);
        self.lock.store(0, Ordering::Relaxed);
    }
}
<commit_msg>Revert "Change `match` to `if let` to satisfy dropck"<commit_after>use core::prelude::*;

use core::atomic::{AtomicUsize, Ordering, ATOMIC_USIZE_INIT};
use core::cell::UnsafeCell;
use core::ops::{Deref, DerefMut};

/// A reader-writer lock
///
/// This type of lock allows a number of readers or at most one writer at any
/// point in time. The write portion of this lock typically allows modification
/// of the underlying data (exclusive access) and the read portion of this lock
/// typically allows for read-only access (shared access).
///
/// The type parameter `T` represents the data that this lock protects. It is
/// required that `T` satisfies `Send` to be shared across tasks and `Sync` to
/// allow concurrent access through readers. The RAII guards returned from the
/// locking methods implement `Deref` (and `DerefMut` for the `write` methods)
/// to allow access to the contained of the lock.
///
/// Based on
/// https://jfdube.wordpress.com/2014/01/03/implementing-a-recursive-read-write-spinlock/
///
pub struct RwLock<T>
{
    lock: AtomicUsize,
    data: UnsafeCell<T>,
}

/// A guard to which the protected data can be read
///
/// When the guard falls out of scope it will decrement the read count,
/// potentially releasing the lock.
pub struct RwLockReadGuard<'a, T:'a>
{
    lock: &'a AtomicUsize,
    data: &'a T,
}

/// A guard to which the protected data can be written
///
/// When the guard falls out of scope it will release the lock.
pub struct RwLockWriteGuard<'a, T:'a>
{
    lock: &'a AtomicUsize,
    data: &'a mut T,
}

unsafe impl<T> Sync for RwLock<T> {}
unsafe impl<T:'static+Send> Send for RwLock<T> {}

/// A RwLock which may be used statically.
///
/// ```
/// use spin::{StaticRwLock, STATIC_RWLOCK_INIT};
///
/// static SPRWLCK: StaticRwLock = STATIC_RWLOCK_INIT;
///
/// fn demo() {
///     let lock = SPRWLCK.read();
///     // do something with lock
///     drop(lock);
/// }
/// ```
pub type StaticRwLock = RwLock<()>;

/// A initializer for StaticRwLock, containing no data.
pub const STATIC_RWLOCK_INIT: StaticRwLock = RwLock {
    lock: ATOMIC_USIZE_INIT,
    data: UnsafeCell { value: () },
};

const USIZE_MSB: usize = ::core::isize::MIN as usize;

impl<T> RwLock<T>
{
    /// Creates a new spinlock wrapping the supplied data.
    #[inline]
    pub fn new(user_data: T) -> RwLock<T>
    {
        RwLock
        {
            lock: ATOMIC_USIZE_INIT,
            data: UnsafeCell::new(user_data),
        }
    }

    /// Locks this rwlock with shared read access, blocking the current thread
    /// until it can be acquired.
    ///
    /// The calling thread will be blocked until there are no more writers which
    /// hold the lock. There may be other readers currently inside the lock when
    /// this method returns. This method does not provide any guarantees with
    /// respect to the ordering of whether contentious readers or writers will
    /// acquire the lock first.
    ///
    /// Returns an RAII guard which will release this thread's shared access
    /// once it is dropped.
    ///
    /// ```
    /// let mylock = spin::RwLock::new(0us);
    /// {
    ///     let mut data = mylock.read();
    ///     // The lock is now locked and the data can be read
    ///     println!("{}", *data);
    ///     // The lock is dropped
    /// }
    /// ```
    #[inline]
    pub fn read<'a>(&'a self) -> RwLockReadGuard<'a, T>
    {
        // (funny do-while loop)
        while {
            // Old value, with write bit unset
            let mut old;

            // Wait for for writer to go away before doing expensive atomic ops
            // (funny do-while loop)
            while {
                old = self.lock.load(Ordering::Relaxed);
                old & USIZE_MSB != 0
            } {}

            // unset write bit
            old &= !USIZE_MSB;

            let new = old + 1;
            debug_assert!(new != (!USIZE_MSB) & (-1));

            self.lock.compare_and_swap(old, new, Ordering::SeqCst) != old
        } {}
        RwLockReadGuard {
            lock: &self.lock,
            data: unsafe { & *self.data.get() },
        }
    }

    /// Attempt to acquire this lock with shared read access.
    ///
    /// This function will never block and will return immediately if `read`
    /// would otherwise succeed. Returns `Some` of an RAII guard which will
    /// release the shared access of this thread when dropped, or `None` if the
    /// access could not be granted. This method does not provide any
    /// guarantees with respect to the ordering of whether contentious readers
    /// or writers will acquire the lock first.
    ///
    /// ```
    /// let mylock = spin::RwLock::new(0us);
    /// {
    ///     match mylock.try_read() {
    ///         Some(data) => {
    ///             // The lock is now locked and the data can be read
    ///             println!("{}", *data);
    ///             // The lock is dropped
    ///         },
    ///         None => (), // no cigar
    ///     }
    /// }
    /// ```
    #[inline]
    pub fn try_read(&self) -> Option<RwLockReadGuard<T>>
    {
        // Old value, with write bit unset
        let old = (!USIZE_MSB) & self.lock.load(Ordering::Relaxed);

        let new = old + 1;
        debug_assert!(new != (!USIZE_MSB) & (-1));
        if self.lock.compare_and_swap(old,
                                      new,
                                      Ordering::SeqCst) == old
        {
            Some(RwLockReadGuard {
                lock: &self.lock,
                data: unsafe { & *self.data.get() },
            })
        } else {
            None
        }
    }

    /// Lock this rwlock with exclusive write access, blocking the current
    /// thread until it can be acquired.
    ///
    /// This function will not return while other writers or other readers
    /// currently have access to the lock.
    ///
    /// Returns an RAII guard which will drop the write access of this rwlock
    /// when dropped.
    ///
    /// ```
    /// let mylock = spin::RwLock::new(0us);
    /// {
    ///     let mut data = mylock.write();
    ///     // The lock is now locked and the data can be written
    ///     *data += 1;
    ///     // The lock is dropped
    /// }
    /// ```
    #[inline]
    pub fn write<'a>(&'a self) -> RwLockWriteGuard<'a, T>
    {
        loop
        {
            // Old value, with write bit unset.
            let old = (!USIZE_MSB) & self.lock.load(Ordering::Relaxed);
            // Old value, with write bit set.
            let new = USIZE_MSB | old;
            if self.lock.compare_and_swap(old,
                                          new,
                                          Ordering::SeqCst) == old
            {
                // Wait for readers to go away, then lock is ours.
                while self.lock.load(Ordering::Relaxed) != USIZE_MSB { }
                break
            }
        }
        RwLockWriteGuard {
            lock: &self.lock,
            data: unsafe { &mut *self.data.get() },
        }
    }

    /// Attempt to lock this rwlock with exclusive write access.
    ///
    /// This function does not ever block, and it will return `None` if a call
    /// to `write` would otherwise block. If successful, an RAII guard is
    /// returned.
    ///
    /// ```
    /// let mylock = spin::RwLock::new(0us);
    /// {
    ///     match mylock.try_write() {
    ///         Some(mut data) => {
    ///             // The lock is now locked and the data can be written
    ///             *data += 1;
    ///             // The lock is implicitly dropped
    ///         },
    ///         None => (), // no cigar
    ///     }
    /// }
    /// ```
    #[inline]
    pub fn try_write(&self) -> Option<RwLockWriteGuard<T>>
    {
        if self.lock.compare_and_swap(0,
                                      USIZE_MSB,
                                      Ordering::SeqCst) == 0
        {
            Some(RwLockWriteGuard {
                lock: &self.lock,
                data: unsafe { &mut *self.data.get() },
            })
        } else {
            None
        }
    }
}

impl<'rwlock, T> Deref for RwLockReadGuard<'rwlock, T> {
    type Target = T;

    fn deref(&self) -> &T { self.data }
}

impl<'rwlock, T> Deref for RwLockWriteGuard<'rwlock, T> {
    type Target = T;

    fn deref(&self) -> &T { self.data }
}

impl<'rwlock, T> DerefMut for RwLockWriteGuard<'rwlock, T> {
    fn deref_mut(&mut self) -> &mut T { self.data }
}

#[unsafe_destructor]
impl<'rwlock, T> Drop for RwLockReadGuard<'rwlock, T> {
    fn drop(&mut self) {
        debug_assert!(self.lock.load(Ordering::Relaxed) & (!USIZE_MSB) > 0);
        self.lock.fetch_sub(1, Ordering::SeqCst);
    }
}

#[unsafe_destructor]
impl<'rwlock, T> Drop for RwLockWriteGuard<'rwlock, T> {
    fn drop(&mut self) {
        debug_assert_eq!(self.lock.load(Ordering::Relaxed), USIZE_MSB);
        self.lock.store(0, Ordering::Relaxed);
    }
}
<|endoftext|>


<commit_before>/*!
Service initialization
*/
use std::env;
use std::thread;
use std::io;

use env_logger;
use chrono::Local;
use rouille;

use handlers;
use sweep;
use db;
use models;
use {ToResponse};
use errors::*;

/// Initialize the `status` database table if it doesn't already exist
fn init_status() -> Result<()> {
    let conn = db::init_conn()?;
    models::Status::init_load(&conn)?;
    Ok(())
}

/// Initialize things
/// - env logger
/// - database `status` table
/// - database connection pool
/// - cleaning thread
/// - server
/// - handle errors
pub fn start(host: &str, port: u16) -> Result<()> {
    // Set a custom logging format & change the env-var to "LOG"
    // e.g. LOG=info chatbot serve
    env_logger::LogBuilder::new()
        .format(|record| {
            format!("{} [{}] - [{}] -> {}",
                Local::now().format("%Y-%m-%d_%H:%M:%S"),
                record.level(),
                record.location().module_path(),
                record.args()
                )
            })
        .parse(&env::var("LOG").unwrap_or_default())
        .init()?;

    // force a config load
    let _ = models::CONFIG;

    // make sure `status` record is initialized
    init_status()?;

    // spawn our cleaning thread
    let _ = thread::spawn(sweep::db_sweeper);

    let db_pool = db::init_pool();

    let addr = format!("{}:{}", host, port);
    info!("** Listening on {} **", addr);

    rouille::start_server(&addr, move |request| {
        let db_pool = db_pool.clone();

        rouille::log(request, io::stdout(), move || {
            // dispatch and handle errors
            match route_request(request, db_pool) {
                Ok(resp) => resp,
                Err(e) => {
                    use self::ErrorKind::*;
                    error!("Handler Error: {}", e);
                    match *e {
                        BadRequest(ref s) => {
                            // bad request
                            let body = json!({"error": s});
                            body.to_resp().unwrap().with_status_code(400)
                        }
                        InvalidAuth(ref s) => {
                            // unauthorized
                            let body = json!({"error": s});
                            body.to_resp().unwrap().with_status_code(401)
                        }
                        DoesNotExist(ref s) => {
                            // not found
                            let body = json!({"error": s});
                            body.to_resp().unwrap().with_status_code(404)
                        }
                        UploadTooLarge(ref s) => {
                            // payload too large / request entity to large
                            let body = json!({"error": s});
                            body.to_resp().unwrap().with_status_code(413)
                        }
                        OutOfSpace(ref s) => {
                            // service unavailable
                            let body = json!({"error": s});
                            body.to_resp().unwrap().with_status_code(503)
                        }
                        _ => rouille::Response::text("Something went wrong").with_status_code(500),
                    }
                }
            }
        })
    });
}

/// Route the request to appropriate handler
fn route_request(request: &rouille::Request, db_pool: db::Pool) -> Result<rouille::Response> {
    Ok(router!(request,
        (GET) (/) => {
            handlers::serve_file("text/html", "assets/main.html")?
        },

        (GET)   (/api/hello)    => { json!({"message": "hey!"}).to_resp()? },
        (POST)  (/api/bye)      => { json!({"message": "bye!"}).to_resp()? },

        (GET)   (/api/upload/defaults)  => { handlers::api_upload_defaults(request)? },
        (POST)  (/api/upload/init)      => { handlers::api_upload_init(request, &db_pool)? },
        (POST)  (/api/upload)           => { handlers::api_upload_file(request, &db_pool)? },
        (POST)  (/api/upload/delete)    => { handlers::api_upload_delete(request, &db_pool)? },

        (POST)  (/api/download/init)    => { handlers::api_download_init(request, &db_pool)? },
        (POST)  (/api/download)         => { handlers::api_download(request, &db_pool)? },
        (POST)  (/api/download/confirm) => { handlers::api_download_confirm(request, &db_pool)? },

        _ => {
            // static files
            let static_resp = rouille::match_assets(&request, "assets");
            if static_resp.is_success() {
                static_resp
            } else {
                bail_fmt!(ErrorKind::DoesNotExist, "nothing here")
            }
        }
    ))
}

<commit_msg>custom request/response logger<commit_after>/*!
Service initialization
*/
use std::env;
use std::thread;
use std::time;

use env_logger;
use chrono::Local;
use rouille;

use handlers;
use sweep;
use db;
use models;
use {ToResponse};
use errors::*;

/// Initialize the `status` database table if it doesn't already exist
fn init_status() -> Result<()> {
    let conn = db::init_conn()?;
    models::Status::init_load(&conn)?;
    Ok(())
}

/// Initialize things
/// - env logger
/// - database `status` table
/// - database connection pool
/// - cleaning thread
/// - server
/// - handle errors
pub fn start(host: &str, port: u16) -> Result<()> {
    // Set a custom logging format & change the env-var to "LOG"
    // e.g. LOG=info chatbot serve
    env_logger::LogBuilder::new()
        .format(|record| {
            format!("{} [{}] - [{}] -> {}",
                Local::now().format("%Y-%m-%d_%H:%M:%S"),
                record.level(),
                record.location().module_path(),
                record.args()
                )
            })
        .parse(&env::var("LOG").unwrap_or_default())
        .init()?;

    // force a config load
    let _ = models::CONFIG;

    // make sure `status` record is initialized
    init_status()?;

    // spawn our cleaning thread
    let _ = thread::spawn(sweep::db_sweeper);

    let db_pool = db::init_pool();

    let addr = format!("{}:{}", host, port);
    info!("** Listening on {} **", addr);

    rouille::start_server(&addr, move |request| {
        let db_pool = db_pool.clone();
        let start = time::Instant::now();

        // dispatch and handle errors
        let response = match route_request(request, db_pool) {
            Ok(resp) => resp,
            Err(e) => {
                use self::ErrorKind::*;
                error!("Handler Error: {}", e);
                match *e {
                    BadRequest(ref s) => {
                        // bad request
                        let body = json!({"error": s});
                        body.to_resp().unwrap().with_status_code(400)
                    }
                    InvalidAuth(ref s) => {
                        // unauthorized
                        let body = json!({"error": s});
                        body.to_resp().unwrap().with_status_code(401)
                    }
                    DoesNotExist(ref s) => {
                        // not found
                        let body = json!({"error": s});
                        body.to_resp().unwrap().with_status_code(404)
                    }
                    UploadTooLarge(ref s) => {
                        // payload too large / request entity to large
                        let body = json!({"error": s});
                        body.to_resp().unwrap().with_status_code(413)
                    }
                    OutOfSpace(ref s) => {
                        // service unavailable
                        let body = json!({"error": s});
                        body.to_resp().unwrap().with_status_code(503)
                    }
                    _ => rouille::Response::text("Something went wrong").with_status_code(500),
                }
            }
        };

        let elapsed = start.elapsed();
        let elapsed = (elapsed.as_secs() * 1_000) as f32 + (elapsed.subsec_nanos() as f32 / 1_000_000.);
        info!("[{}] {} {:?} {}ms", request.method(), response.status_code, request.url(), elapsed);
        response
    });
}

/// Route the request to appropriate handler
fn route_request(request: &rouille::Request, db_pool: db::Pool) -> Result<rouille::Response> {
    Ok(router!(request,
        (GET) (/) => {
            handlers::serve_file("text/html", "assets/main.html")?
        },

        (GET)   (/api/hello)    => { json!({"message": "hey!"}).to_resp()? },
        (POST)  (/api/bye)      => { json!({"message": "bye!"}).to_resp()? },

        (GET)   (/api/upload/defaults)  => { handlers::api_upload_defaults(request)? },
        (POST)  (/api/upload/init)      => { handlers::api_upload_init(request, &db_pool)? },
        (POST)  (/api/upload)           => { handlers::api_upload_file(request, &db_pool)? },
        (POST)  (/api/upload/delete)    => { handlers::api_upload_delete(request, &db_pool)? },

        (POST)  (/api/download/init)    => { handlers::api_download_init(request, &db_pool)? },
        (POST)  (/api/download)         => { handlers::api_download(request, &db_pool)? },
        (POST)  (/api/download/confirm) => { handlers::api_download_confirm(request, &db_pool)? },

        _ => {
            // static files
            let static_resp = rouille::match_assets(&request, "assets");
            if static_resp.is_success() {
                static_resp
            } else {
                bail_fmt!(ErrorKind::DoesNotExist, "nothing here")
            }
        }
    ))
}

<|endoftext|>


<commit_before>mod helpers;

use helpers::{roundtrip, wayc, ways, TestClient, TestServer};

use ways::protocol::{wl_compositor, wl_output};

use wayc::protocol::wl_compositor::WlCompositor as ClientCompositor;
use wayc::protocol::wl_output::WlOutput as ClientOutput;

use std::sync::{Arc, Mutex};

#[test]
fn client_user_data() {
    let mut server = TestServer::new();
    let clients = Arc::new(Mutex::new(Vec::new()));

    struct HasOutput;
    struct HasCompositor;

    server.display.create_global::<wl_output::WlOutput, _>(1, {
        let clients = clients.clone();
        ways::Filter::new(move |(output, _): (ways::Main<wl_output::WlOutput>, u32), _, _| {
            let client = output.as_ref().client().unwrap();
            let ret = client.data_map().insert_if_missing(|| HasOutput);
            // the data should not be already here
            assert!(ret);
            clients.lock().unwrap().push(client);
        })
    });
    server.display.create_global::<wl_compositor::WlCompositor, _>(1, {
        let clients = clients.clone();
        ways::Filter::new(
            move |(compositor, _): (ways::Main<wl_compositor::WlCompositor>, u32), _, _| {
                let client = compositor.as_ref().client().unwrap();
                let ret = client.data_map().insert_if_missing(|| HasCompositor);
                // the data should not be already here
                assert!(ret);
                clients.lock().unwrap().push(client);
            },
        )
    });

    let mut client = TestClient::new(&server.socket_name);
    let manager = wayc::GlobalManager::new(&client.display_proxy);

    roundtrip(&mut client, &mut server).unwrap();

    // Instantiate the globals
    manager.instantiate_exact::<ClientOutput>(1).unwrap();

    roundtrip(&mut client, &mut server).unwrap();

    {
        let clients = clients.lock().unwrap();

        assert!(clients.len() == 1);
        assert!(clients[0].data_map().get::<HasOutput>().is_some());
        assert!(clients[0].data_map().get::<HasCompositor>().is_none());
    }

    manager.instantiate_exact::<ClientCompositor>(1).unwrap();

    roundtrip(&mut client, &mut server).unwrap();

    let clients = clients.lock().unwrap();

    assert!(clients.len() == 2);
    assert!(clients[0].equals(&clients[1]));
    assert!(clients[0].data_map().get::<HasCompositor>().is_some());
    assert!(clients[0].data_map().get::<HasOutput>().is_some());
    assert!(clients[1].data_map().get::<HasCompositor>().is_some());
    assert!(clients[1].data_map().get::<HasOutput>().is_some());
}

#[test]
fn client_credentials() {
    let mut server = TestServer::new();

    let server_client = Arc::new(Mutex::new(None));

    server.display.create_global::<wl_output::WlOutput, _>(1, {
        let server_client = server_client.clone();
        ways::Filter::new(move |(output, _): (ways::Main<wl_output::WlOutput>, u32), _, _| {
            let client = output.as_ref().client().unwrap();
            server_client.lock().unwrap().replace(client);
        })
    });
    let mut client = TestClient::new(&server.socket_name);
    let manager = wayc::GlobalManager::new(&client.display_proxy);

    roundtrip(&mut client, &mut server).unwrap();

    // Instantiate the globals
    manager.instantiate_exact::<ClientOutput>(1).unwrap();

    roundtrip(&mut client, &mut server).unwrap();

    let server_client = server_client.lock().unwrap().take().unwrap();
    let credentials = server_client.credentials();
    assert!(credentials.is_some());

    let credentials = credentials.unwrap();

    assert!(credentials.pid != 0);
}
<commit_msg>ignore client_credentials test assert on freebsd<commit_after>mod helpers;

use helpers::{roundtrip, wayc, ways, TestClient, TestServer};

use ways::protocol::{wl_compositor, wl_output};

use wayc::protocol::wl_compositor::WlCompositor as ClientCompositor;
use wayc::protocol::wl_output::WlOutput as ClientOutput;

use std::sync::{Arc, Mutex};

#[test]
fn client_user_data() {
    let mut server = TestServer::new();
    let clients = Arc::new(Mutex::new(Vec::new()));

    struct HasOutput;
    struct HasCompositor;

    server.display.create_global::<wl_output::WlOutput, _>(1, {
        let clients = clients.clone();
        ways::Filter::new(move |(output, _): (ways::Main<wl_output::WlOutput>, u32), _, _| {
            let client = output.as_ref().client().unwrap();
            let ret = client.data_map().insert_if_missing(|| HasOutput);
            // the data should not be already here
            assert!(ret);
            clients.lock().unwrap().push(client);
        })
    });
    server.display.create_global::<wl_compositor::WlCompositor, _>(1, {
        let clients = clients.clone();
        ways::Filter::new(
            move |(compositor, _): (ways::Main<wl_compositor::WlCompositor>, u32), _, _| {
                let client = compositor.as_ref().client().unwrap();
                let ret = client.data_map().insert_if_missing(|| HasCompositor);
                // the data should not be already here
                assert!(ret);
                clients.lock().unwrap().push(client);
            },
        )
    });

    let mut client = TestClient::new(&server.socket_name);
    let manager = wayc::GlobalManager::new(&client.display_proxy);

    roundtrip(&mut client, &mut server).unwrap();

    // Instantiate the globals
    manager.instantiate_exact::<ClientOutput>(1).unwrap();

    roundtrip(&mut client, &mut server).unwrap();

    {
        let clients = clients.lock().unwrap();

        assert!(clients.len() == 1);
        assert!(clients[0].data_map().get::<HasOutput>().is_some());
        assert!(clients[0].data_map().get::<HasCompositor>().is_none());
    }

    manager.instantiate_exact::<ClientCompositor>(1).unwrap();

    roundtrip(&mut client, &mut server).unwrap();

    let clients = clients.lock().unwrap();

    assert!(clients.len() == 2);
    assert!(clients[0].equals(&clients[1]));
    assert!(clients[0].data_map().get::<HasCompositor>().is_some());
    assert!(clients[0].data_map().get::<HasOutput>().is_some());
    assert!(clients[1].data_map().get::<HasCompositor>().is_some());
    assert!(clients[1].data_map().get::<HasOutput>().is_some());
}

#[test]
fn client_credentials() {
    let mut server = TestServer::new();

    let server_client = Arc::new(Mutex::new(None));

    server.display.create_global::<wl_output::WlOutput, _>(1, {
        let server_client = server_client.clone();
        ways::Filter::new(move |(output, _): (ways::Main<wl_output::WlOutput>, u32), _, _| {
            let client = output.as_ref().client().unwrap();
            server_client.lock().unwrap().replace(client);
        })
    });
    let mut client = TestClient::new(&server.socket_name);
    let manager = wayc::GlobalManager::new(&client.display_proxy);

    roundtrip(&mut client, &mut server).unwrap();

    // Instantiate the globals
    manager.instantiate_exact::<ClientOutput>(1).unwrap();

    roundtrip(&mut client, &mut server).unwrap();

    let server_client = server_client.lock().unwrap().take().unwrap();
    let credentials = server_client.credentials();
    assert!(credentials.is_some());
    assert_credentials(credentials.unwrap());
}

#[cfg(any(not(feature = "server_native"), not(target_os = "freebsd")))]
fn assert_credentials(credentials: ways::Credentials) {
    assert!(credentials.pid != 0);
}

#[cfg(all(feature = "server_native", target_os = "freebsd"))]
fn assert_credentials(_credentials: ways::Credentials) {
    // The current implementation of wl_client_get_credentials
    // will always return pid == 0 on freebsd
    // On recent versions this has been fixed with a freebsd
    // specific patch. Detecting if a patched version is used
    // is too complicated and this assert would just test the
    // native wayland-server library. So the assert is a no-op
    // for now.
    //
    // see: https://bugs.freebsd.org/bugzilla/show_bug.cgi?id=246189
}
<|endoftext|>


<commit_before>use crate::ops::{Deref, DerefMut};
use crate::ptr;

/// A wrapper to inhibit compiler from automatically calling `T`’s destructor.
///
/// This wrapper is 0-cost.
///
/// `ManuallyDrop<T>` is subject to the same layout optimizations as `T`.
/// As a consequence, it has *no effect* on the assumptions that the compiler makes
/// about all values being initialized at their type.  In particular, initializing
/// a `ManuallyDrop<&mut T>` with [`mem::zeroed`] is undefined behavior.
/// If you need to handle uninitialized data, use [`MaybeUninit<T>`] instead.
///
/// # Examples
///
/// This wrapper helps with explicitly documenting the drop order dependencies between fields of
/// the type:
///
/// ```rust
/// use std::mem::ManuallyDrop;
/// struct Peach;
/// struct Banana;
/// struct Melon;
/// struct FruitBox {
///     // Immediately clear there’s something non-trivial going on with these fields.
///     peach: ManuallyDrop<Peach>,
///     melon: Melon, // Field that’s independent of the other two.
///     banana: ManuallyDrop<Banana>,
/// }
///
/// impl Drop for FruitBox {
///     fn drop(&mut self) {
///         unsafe {
///             // Explicit ordering in which field destructors are run specified in the intuitive
///             // location – the destructor of the structure containing the fields.
///             // Moreover, one can now reorder fields within the struct however much they want.
///             ManuallyDrop::drop(&mut self.peach);
///             ManuallyDrop::drop(&mut self.banana);
///         }
///         // After destructor for `FruitBox` runs (this function), the destructor for Melon gets
///         // invoked in the usual manner, as it is not wrapped in `ManuallyDrop`.
///     }
/// }
/// ```
///
/// [`mem::zeroed`]: fn.zeroed.html
/// [`MaybeUninit<T>`]: union.MaybeUninit.html
#[stable(feature = "manually_drop", since = "1.20.0")]
#[lang = "manually_drop"]
#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(transparent)]
pub struct ManuallyDrop<T: ?Sized> {
    value: T,
}

impl<T> ManuallyDrop<T> {
    /// Wrap a value to be manually dropped.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use std::mem::ManuallyDrop;
    /// ManuallyDrop::new(Box::new(()));
    /// ```
    #[stable(feature = "manually_drop", since = "1.20.0")]
    #[rustc_const_stable(feature = "const_manually_drop", since = "1.36.0")]
    #[inline(always)]
    pub const fn new(value: T) -> ManuallyDrop<T> {
        ManuallyDrop { value }
    }

    /// Extracts the value from the `ManuallyDrop` container.
    ///
    /// This allows the value to be dropped again.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use std::mem::ManuallyDrop;
    /// let x = ManuallyDrop::new(Box::new(()));
    /// let _: Box<()> = ManuallyDrop::into_inner(x); // This drops the `Box`.
    /// ```
    #[stable(feature = "manually_drop", since = "1.20.0")]
    #[rustc_const_stable(feature = "const_manually_drop", since = "1.36.0")]
    #[inline(always)]
    pub const fn into_inner(slot: ManuallyDrop<T>) -> T {
        slot.value
    }

    /// Takes the contained value out.
    ///
    /// This method is primarily intended for moving out values in drop.
    /// Instead of using [`ManuallyDrop::drop`] to manually drop the value,
    /// you can use this method to take the value and use it however desired.
    /// `Drop` will be invoked on the returned value following normal end-of-scope rules.
    ///
    /// If you have ownership of the container, you can use [`ManuallyDrop::into_inner`] instead.
    ///
    /// # Safety
    ///
    /// This function semantically moves out the contained value without preventing further usage.
    /// It is up to the user of this method to ensure that this container is not used again.
    ///
    /// [`ManuallyDrop::drop`]: #method.drop
    /// [`ManuallyDrop::into_inner`]: #method.into_inner
    #[must_use = "if you don't need the value, you can use `ManuallyDrop::drop` instead"]
    #[unstable(feature = "manually_drop_take", issue = "55422")]
    #[inline]
    pub unsafe fn take(slot: &mut ManuallyDrop<T>) -> T {
        ManuallyDrop::into_inner(ptr::read(slot))
    }
}

impl<T: ?Sized> ManuallyDrop<T> {
    /// Manually drops the contained value.
    ///
    /// If you have ownership of the value, you can use [`ManuallyDrop::into_inner`] instead.
    ///
    /// # Safety
    ///
    /// This function runs the destructor of the contained value and thus the wrapped value
    /// now represents uninitialized data. It is up to the user of this method to ensure the
    /// uninitialized data is not actually used.
    /// In particular, this function can only be called called at most once
    /// for a given instance of `ManuallyDrop<T>`.
    ///
    /// [`ManuallyDrop::into_inner`]: #method.into_inner
    #[stable(feature = "manually_drop", since = "1.20.0")]
    #[inline]
    pub unsafe fn drop(slot: &mut ManuallyDrop<T>) {
        ptr::drop_in_place(&mut slot.value)
    }
}

#[stable(feature = "manually_drop", since = "1.20.0")]
impl<T: ?Sized> Deref for ManuallyDrop<T> {
    type Target = T;
    #[inline(always)]
    fn deref(&self) -> &T {
        &self.value
    }
}

#[stable(feature = "manually_drop", since = "1.20.0")]
impl<T: ?Sized> DerefMut for ManuallyDrop<T> {
    #[inline(always)]
    fn deref_mut(&mut self) -> &mut T {
        &mut self.value
    }
}
<commit_msg>Rollup merge of #67715 - petertodd:2019-typo-manuallydrop, r=Centril<commit_after>use crate::ops::{Deref, DerefMut};
use crate::ptr;

/// A wrapper to inhibit compiler from automatically calling `T`’s destructor.
///
/// This wrapper is 0-cost.
///
/// `ManuallyDrop<T>` is subject to the same layout optimizations as `T`.
/// As a consequence, it has *no effect* on the assumptions that the compiler makes
/// about all values being initialized at their type.  In particular, initializing
/// a `ManuallyDrop<&mut T>` with [`mem::zeroed`] is undefined behavior.
/// If you need to handle uninitialized data, use [`MaybeUninit<T>`] instead.
///
/// # Examples
///
/// This wrapper helps with explicitly documenting the drop order dependencies between fields of
/// the type:
///
/// ```rust
/// use std::mem::ManuallyDrop;
/// struct Peach;
/// struct Banana;
/// struct Melon;
/// struct FruitBox {
///     // Immediately clear there’s something non-trivial going on with these fields.
///     peach: ManuallyDrop<Peach>,
///     melon: Melon, // Field that’s independent of the other two.
///     banana: ManuallyDrop<Banana>,
/// }
///
/// impl Drop for FruitBox {
///     fn drop(&mut self) {
///         unsafe {
///             // Explicit ordering in which field destructors are run specified in the intuitive
///             // location – the destructor of the structure containing the fields.
///             // Moreover, one can now reorder fields within the struct however much they want.
///             ManuallyDrop::drop(&mut self.peach);
///             ManuallyDrop::drop(&mut self.banana);
///         }
///         // After destructor for `FruitBox` runs (this function), the destructor for Melon gets
///         // invoked in the usual manner, as it is not wrapped in `ManuallyDrop`.
///     }
/// }
/// ```
///
/// [`mem::zeroed`]: fn.zeroed.html
/// [`MaybeUninit<T>`]: union.MaybeUninit.html
#[stable(feature = "manually_drop", since = "1.20.0")]
#[lang = "manually_drop"]
#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(transparent)]
pub struct ManuallyDrop<T: ?Sized> {
    value: T,
}

impl<T> ManuallyDrop<T> {
    /// Wrap a value to be manually dropped.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use std::mem::ManuallyDrop;
    /// ManuallyDrop::new(Box::new(()));
    /// ```
    #[stable(feature = "manually_drop", since = "1.20.0")]
    #[rustc_const_stable(feature = "const_manually_drop", since = "1.36.0")]
    #[inline(always)]
    pub const fn new(value: T) -> ManuallyDrop<T> {
        ManuallyDrop { value }
    }

    /// Extracts the value from the `ManuallyDrop` container.
    ///
    /// This allows the value to be dropped again.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use std::mem::ManuallyDrop;
    /// let x = ManuallyDrop::new(Box::new(()));
    /// let _: Box<()> = ManuallyDrop::into_inner(x); // This drops the `Box`.
    /// ```
    #[stable(feature = "manually_drop", since = "1.20.0")]
    #[rustc_const_stable(feature = "const_manually_drop", since = "1.36.0")]
    #[inline(always)]
    pub const fn into_inner(slot: ManuallyDrop<T>) -> T {
        slot.value
    }

    /// Takes the contained value out.
    ///
    /// This method is primarily intended for moving out values in drop.
    /// Instead of using [`ManuallyDrop::drop`] to manually drop the value,
    /// you can use this method to take the value and use it however desired.
    /// `Drop` will be invoked on the returned value following normal end-of-scope rules.
    ///
    /// If you have ownership of the container, you can use [`ManuallyDrop::into_inner`] instead.
    ///
    /// # Safety
    ///
    /// This function semantically moves out the contained value without preventing further usage.
    /// It is up to the user of this method to ensure that this container is not used again.
    ///
    /// [`ManuallyDrop::drop`]: #method.drop
    /// [`ManuallyDrop::into_inner`]: #method.into_inner
    #[must_use = "if you don't need the value, you can use `ManuallyDrop::drop` instead"]
    #[unstable(feature = "manually_drop_take", issue = "55422")]
    #[inline]
    pub unsafe fn take(slot: &mut ManuallyDrop<T>) -> T {
        ManuallyDrop::into_inner(ptr::read(slot))
    }
}

impl<T: ?Sized> ManuallyDrop<T> {
    /// Manually drops the contained value.
    ///
    /// If you have ownership of the value, you can use [`ManuallyDrop::into_inner`] instead.
    ///
    /// # Safety
    ///
    /// This function runs the destructor of the contained value and thus the wrapped value
    /// now represents uninitialized data. It is up to the user of this method to ensure the
    /// uninitialized data is not actually used.
    /// In particular, this function can only be called at most once
    /// for a given instance of `ManuallyDrop<T>`.
    ///
    /// [`ManuallyDrop::into_inner`]: #method.into_inner
    #[stable(feature = "manually_drop", since = "1.20.0")]
    #[inline]
    pub unsafe fn drop(slot: &mut ManuallyDrop<T>) {
        ptr::drop_in_place(&mut slot.value)
    }
}

#[stable(feature = "manually_drop", since = "1.20.0")]
impl<T: ?Sized> Deref for ManuallyDrop<T> {
    type Target = T;
    #[inline(always)]
    fn deref(&self) -> &T {
        &self.value
    }
}

#[stable(feature = "manually_drop", since = "1.20.0")]
impl<T: ?Sized> DerefMut for ManuallyDrop<T> {
    #[inline(always)]
    fn deref_mut(&mut self) -> &mut T {
        &mut self.value
    }
}
<|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.

//! Sampling from random distributions

// Some implementations use the Ziggurat method
// https://en.wikipedia.org/wiki/Ziggurat_algorithm
//
// The version used here is ZIGNOR [Doornik 2005, "An Improved
// Ziggurat Method to Generate Normal Random Samples"] which is slower
// (about double, it generates an extra random number) than the
// canonical version [Marsaglia & Tsang 2000, "The Ziggurat Method for
// Generating Random Variables"], but more robust. If one wanted, one
// could implement VIZIGNOR the ZIGNOR paper for more speed.

use num;
use rand::{Rng,Rand};

/// Things that can be used to create a random instance of `Support`.
pub trait Sample<Support> {
    /// Generate a random value of `Support`, using `rng` as the
    /// source of randomness.
    fn sample<R: Rng>(&mut self, rng: &mut R) -> Support;
}

/// `Sample`s that do not require keeping track of state, so each
/// sample is (statistically) independent of all others, assuming the
/// `Rng` used has this property.
// XXX maybe having this separate is overkill (the only reason is to
// take &self rather than &mut self)? or maybe this should be the
// trait called `Sample` and the other should be `DependentSample`.
pub trait IndependentSample<Support>: Sample<Support> {
    /// Generate a random value.
    fn ind_sample<R: Rng>(&self, &mut R) -> Support;
}

/// A wrapper for generating types that implement `Rand` via the
/// `Sample` & `IndependentSample` traits.
pub struct RandSample<Sup>;

impl<Sup: Rand> Sample<Sup> for RandSample<Sup> {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> Sup { self.ind_sample(rng) }
}

impl<Sup: Rand> IndependentSample<Sup> for RandSample<Sup> {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> Sup {
        rng.gen()
    }
}

mod ziggurat_tables;

// inlining should mean there is no performance penalty for this
#[inline]
fn ziggurat<R:Rng>(rng: &mut R,
                   center_u: bool,
                   X: ziggurat_tables::ZigTable,
                   F: ziggurat_tables::ZigTable,
                   F_DIFF: ziggurat_tables::ZigTable,
                   pdf: &'static fn(f64) -> f64, // probability density function
                   zero_case: &'static fn(&mut R, f64) -> f64) -> f64 {
    loop {
        let u = if center_u {2.0 * rng.gen() - 1.0} else {rng.gen()};
        let i: uint = rng.gen::<uint>() & 0xff;
        let x = u * X[i];

        let test_x = if center_u {num::abs(x)} else {x};

        // algebraically equivalent to |u| < X[i+1]/X[i] (or u < X[i+1]/X[i])
        if test_x < X[i + 1] {
            return x;
        }
        if i == 0 {
            return zero_case(rng, u);
        }
        // algebraically equivalent to f1 + DRanU()*(f0 - f1) < 1
        if F[i+1] + F_DIFF[i+1] * rng.gen() < pdf(x) {
            return x;
        }
    }
}

/// A wrapper around an `f64` to generate N(0, 1) random numbers (a.k.a.  a
/// standard normal, or Gaussian). Multiplying the generated values by the
/// desired standard deviation `sigma` then adding the desired mean `mu` will
/// give N(mu, sigma^2) distributed random numbers.
///
/// Note that this has to be unwrapped before use as an `f64` (using either
/// `*` or `cast::transmute` is safe).
///
/// # Example
///
/// ```
/// use std::rand::distributions::StandardNormal;
///
/// fn main() {
///     let normal = 2.0 + (*rand::random::<StandardNormal>()) * 3.0;
///     println!("{} is from a N(2, 9) distribution", normal)
/// }
/// ```
pub struct StandardNormal(f64);

impl Rand for StandardNormal {
    fn rand<R:Rng>(rng: &mut R) -> StandardNormal {
        #[inline]
        fn pdf(x: f64) -> f64 {
            ((-x*x/2.0) as f64).exp()
        }
        #[inline]
        fn zero_case<R:Rng>(rng: &mut R, u: f64) -> f64 {
            // compute a random number in the tail by hand

            // strange initial conditions, because the loop is not
            // do-while, so the condition should be true on the first
            // run, they get overwritten anyway (0 < 1, so these are
            // good).
            let mut x = 1.0f64;
            let mut y = 0.0f64;

            // FIXME #7755: infinities?
            while -2.0 * y < x * x {
                x = rng.gen::<f64>().ln() / ziggurat_tables::ZIG_NORM_R;
                y = rng.gen::<f64>().ln();
            }

            if u < 0.0 { x - ziggurat_tables::ZIG_NORM_R } else { ziggurat_tables::ZIG_NORM_R - x }
        }

        StandardNormal(ziggurat(
            rng,
            true, // this is symmetric
            &ziggurat_tables::ZIG_NORM_X,
            &ziggurat_tables::ZIG_NORM_F, &ziggurat_tables::ZIG_NORM_F_DIFF,
            pdf, zero_case))
    }
}

/// A wrapper around an `f64` to generate Exp(1) random numbers. Dividing by
/// the desired rate `lambda` will give Exp(lambda) distributed random
/// numbers.
///
/// Note that this has to be unwrapped before use as an `f64` (using either
/// `*` or `cast::transmute` is safe).
///
/// # Example
///
/// ```
/// use std::rand::distributions::Exp1;
///
/// fn main() {
///     let exp2 = (*rand::random::<Exp1>()) * 0.5;
///     println!("{} is from a Exp(2) distribution", exp2);
/// }
/// ```
pub struct Exp1(f64);

// This could be done via `-rng.gen::<f64>().ln()` but that is slower.
impl Rand for Exp1 {
    #[inline]
    fn rand<R:Rng>(rng: &mut R) -> Exp1 {
        #[inline]
        fn pdf(x: f64) -> f64 {
            (-x).exp()
        }
        #[inline]
        fn zero_case<R:Rng>(rng: &mut R, _u: f64) -> f64 {
            ziggurat_tables::ZIG_EXP_R - rng.gen::<f64>().ln()
        }

        Exp1(ziggurat(rng, false,
                      &ziggurat_tables::ZIG_EXP_X,
                      &ziggurat_tables::ZIG_EXP_F, &ziggurat_tables::ZIG_EXP_F_DIFF,
                      pdf, zero_case))
    }
}

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

    struct ConstRand(uint);
    impl Rand for ConstRand {
        fn rand<R: Rng>(_: &mut R) -> ConstRand {
            ConstRand(0)
        }
    }

    #[test]
    fn test_rand_sample() {
        let mut rand_sample = RandSample::<ConstRand>;

        assert_eq!(*rand_sample.sample(task_rng()), 0);
        assert_eq!(*rand_sample.ind_sample(task_rng()), 0);
    }
}
<commit_msg>std::rand: full exponential & normal distributions<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.

//! Sampling from random distributions

// Some implementations use the Ziggurat method
// https://en.wikipedia.org/wiki/Ziggurat_algorithm
//
// The version used here is ZIGNOR [Doornik 2005, "An Improved
// Ziggurat Method to Generate Normal Random Samples"] which is slower
// (about double, it generates an extra random number) than the
// canonical version [Marsaglia & Tsang 2000, "The Ziggurat Method for
// Generating Random Variables"], but more robust. If one wanted, one
// could implement VIZIGNOR the ZIGNOR paper for more speed.

use num;
use rand::{Rng,Rand};

/// Things that can be used to create a random instance of `Support`.
pub trait Sample<Support> {
    /// Generate a random value of `Support`, using `rng` as the
    /// source of randomness.
    fn sample<R: Rng>(&mut self, rng: &mut R) -> Support;
}

/// `Sample`s that do not require keeping track of state, so each
/// sample is (statistically) independent of all others, assuming the
/// `Rng` used has this property.
// XXX maybe having this separate is overkill (the only reason is to
// take &self rather than &mut self)? or maybe this should be the
// trait called `Sample` and the other should be `DependentSample`.
pub trait IndependentSample<Support>: Sample<Support> {
    /// Generate a random value.
    fn ind_sample<R: Rng>(&self, &mut R) -> Support;
}

/// A wrapper for generating types that implement `Rand` via the
/// `Sample` & `IndependentSample` traits.
pub struct RandSample<Sup>;

impl<Sup: Rand> Sample<Sup> for RandSample<Sup> {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> Sup { self.ind_sample(rng) }
}

impl<Sup: Rand> IndependentSample<Sup> for RandSample<Sup> {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> Sup {
        rng.gen()
    }
}

mod ziggurat_tables;

// inlining should mean there is no performance penalty for this
#[inline]
fn ziggurat<R:Rng>(rng: &mut R,
                   center_u: bool,
                   X: ziggurat_tables::ZigTable,
                   F: ziggurat_tables::ZigTable,
                   F_DIFF: ziggurat_tables::ZigTable,
                   pdf: &'static fn(f64) -> f64, // probability density function
                   zero_case: &'static fn(&mut R, f64) -> f64) -> f64 {
    loop {
        let u = if center_u {2.0 * rng.gen() - 1.0} else {rng.gen()};
        let i: uint = rng.gen::<uint>() & 0xff;
        let x = u * X[i];

        let test_x = if center_u {num::abs(x)} else {x};

        // algebraically equivalent to |u| < X[i+1]/X[i] (or u < X[i+1]/X[i])
        if test_x < X[i + 1] {
            return x;
        }
        if i == 0 {
            return zero_case(rng, u);
        }
        // algebraically equivalent to f1 + DRanU()*(f0 - f1) < 1
        if F[i+1] + F_DIFF[i+1] * rng.gen() < pdf(x) {
            return x;
        }
    }
}

/// A wrapper around an `f64` to generate N(0, 1) random numbers (a.k.a.  a
/// standard normal, or Gaussian). Multiplying the generated values by the
/// desired standard deviation `sigma` then adding the desired mean `mu` will
/// give N(mu, sigma^2) distributed random numbers.
///
/// Note that this has to be unwrapped before use as an `f64` (using either
/// `*` or `cast::transmute` is safe).
pub struct StandardNormal(f64);

impl Rand for StandardNormal {
    fn rand<R:Rng>(rng: &mut R) -> StandardNormal {
        #[inline]
        fn pdf(x: f64) -> f64 {
            ((-x*x/2.0) as f64).exp()
        }
        #[inline]
        fn zero_case<R:Rng>(rng: &mut R, u: f64) -> f64 {
            // compute a random number in the tail by hand

            // strange initial conditions, because the loop is not
            // do-while, so the condition should be true on the first
            // run, they get overwritten anyway (0 < 1, so these are
            // good).
            let mut x = 1.0f64;
            let mut y = 0.0f64;

            // FIXME #7755: infinities?
            while -2.0 * y < x * x {
                x = rng.gen::<f64>().ln() / ziggurat_tables::ZIG_NORM_R;
                y = rng.gen::<f64>().ln();
            }

            if u < 0.0 { x - ziggurat_tables::ZIG_NORM_R } else { ziggurat_tables::ZIG_NORM_R - x }
        }

        StandardNormal(ziggurat(
            rng,
            true, // this is symmetric
            &ziggurat_tables::ZIG_NORM_X,
            &ziggurat_tables::ZIG_NORM_F, &ziggurat_tables::ZIG_NORM_F_DIFF,
            pdf, zero_case))
    }
}

/// The `N(mean, std_dev**2)` distribution, i.e. samples from a normal
/// distribution with mean `mean` and standard deviation `std_dev`.
///
/// # Example
///
/// ```
/// use std::rand;
/// use std::rand::distributions::{Normal, IndependentSample};
///
/// fn main() {
///     let normal = Normal::new(2.0, 3.0);
///     let v = normal.ind_sample(rand::task_rng());
///     println!("{} is from a N(2, 9) distribution", v)
/// }
/// ```
pub struct Normal {
    priv mean: f64,
    priv std_dev: f64
}

impl Normal {
    /// Construct a new `Normal` distribution with the given mean and
    /// standard deviation. Fails if `std_dev < 0`.
    pub fn new(mean: f64, std_dev: f64) -> Normal {
        assert!(std_dev >= 0.0, "Normal::new called with `std_dev` < 0");
        Normal {
            mean: mean,
            std_dev: std_dev
        }
    }
}
impl Sample<f64> for Normal {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl IndependentSample<f64> for Normal {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        self.mean + self.std_dev * (*rng.gen::<StandardNormal>())
    }
}

/// A wrapper around an `f64` to generate Exp(1) random numbers. Dividing by
/// the desired rate `lambda` will give Exp(lambda) distributed random
/// numbers.
///
/// Note that this has to be unwrapped before use as an `f64` (using either
/// `*` or `cast::transmute` is safe).
pub struct Exp1(f64);

// This could be done via `-rng.gen::<f64>().ln()` but that is slower.
impl Rand for Exp1 {
    #[inline]
    fn rand<R:Rng>(rng: &mut R) -> Exp1 {
        #[inline]
        fn pdf(x: f64) -> f64 {
            (-x).exp()
        }
        #[inline]
        fn zero_case<R:Rng>(rng: &mut R, _u: f64) -> f64 {
            ziggurat_tables::ZIG_EXP_R - rng.gen::<f64>().ln()
        }

        Exp1(ziggurat(rng, false,
                      &ziggurat_tables::ZIG_EXP_X,
                      &ziggurat_tables::ZIG_EXP_F, &ziggurat_tables::ZIG_EXP_F_DIFF,
                      pdf, zero_case))
    }
}

/// The `Exp(lambda)` distribution; i.e. samples from the exponential
/// distribution with rate parameter `lambda`.
///
/// This distribution has density function: `f(x) = lambda *
/// exp(-lambda * x)` for `x > 0`.
///
/// # Example
///
/// ```
/// use std::rand;
/// use std::rand::distributions::{Exp, IndependentSample};
///
/// fn main() {
///     let exp = Exp::new(2.0);
///     let v = exp.ind_sample(rand::task_rng());
///     println!("{} is from a Exp(2) distribution", v);
/// }
/// ```
pub struct Exp {
    /// `lambda` stored as `1/lambda`, since this is what we scale by.
    priv lambda_inverse: f64
}

impl Exp {
    /// Construct a new `Exp` with the given shape parameter
    /// `lambda`. Fails if `lambda <= 0`.
    pub fn new(lambda: f64) -> Exp {
        assert!(lambda > 0.0, "Exp::new called with `lambda` <= 0");
        Exp { lambda_inverse: 1.0 / lambda }
    }
}

impl Sample<f64> for Exp {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl IndependentSample<f64> for Exp {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        (*rng.gen::<Exp1>()) * self.lambda_inverse
    }
}

#[cfg(test)]
mod tests {
    use rand::*;
    use super::*;
    use iter::range;
    use option::{Some, None};

    struct ConstRand(uint);
    impl Rand for ConstRand {
        fn rand<R: Rng>(_: &mut R) -> ConstRand {
            ConstRand(0)
        }
    }

    #[test]
    fn test_rand_sample() {
        let mut rand_sample = RandSample::<ConstRand>;

        assert_eq!(*rand_sample.sample(task_rng()), 0);
        assert_eq!(*rand_sample.ind_sample(task_rng()), 0);
    }

    #[test]
    fn test_normal() {
        let mut norm = Normal::new(10.0, 10.0);
        let rng = task_rng();
        for _ in range(0, 1000) {
            norm.sample(rng);
            norm.ind_sample(rng);
        }
    }
    #[test]
    #[should_fail]
    fn test_normal_invalid_sd() {
        Normal::new(10.0, -1.0);
    }

    #[test]
    fn test_exp() {
        let mut exp = Exp::new(10.0);
        let rng = task_rng();
        for _ in range(0, 1000) {
            assert!(exp.sample(rng) >= 0.0);
            assert!(exp.ind_sample(rng) >= 0.0);
        }
    }
    #[test]
    #[should_fail]
    fn test_exp_invalid_lambda_zero() {
        Exp::new(0.0);
    }
    #[test]
    #[should_fail]
    fn test_exp_invalid_lambda_neg() {
        Exp::new(-10.0);
    }
}
<|endoftext|>