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

//! The implementations of the `Uniform` distribution for other built-in types.

use core::char;

use {Rng};
use distributions::{Distribution, Uniform, Range};

impl Distribution<char> for Uniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> char {
        let range = Range::new(0u32, 0x11_0000);
        loop {
            match char::from_u32(range.sample(rng)) {
                Some(c) => return c,
                // About 0.2% of numbers in the range 0..0x110000 are invalid
                // codepoints (surrogates).
                None => {}
            }
        }
    }
}

impl Distribution<bool> for Uniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> bool {
        rng.gen::<u8>() & 1 == 1
    }
}

macro_rules! tuple_impl {
    // use variables to indicate the arity of the tuple
    ($($tyvar:ident),* ) => {
        // the trailing commas are for the 1 tuple
        impl< $( $tyvar ),* >
            Distribution<( $( $tyvar ),* , )>
            for Uniform
            where $( Uniform: Distribution<$tyvar> ),*
        {
            #[inline]
            fn sample<R: Rng + ?Sized>(&self, _rng: &mut R) -> ( $( $tyvar ),* , ) {
                (
                    // use the $tyvar's to get the appropriate number of
                    // repeats (they're not actually needed)
                    $(
                        _rng.gen::<$tyvar>()
                    ),*
                    ,
                )
            }
        }
    }
}

impl Distribution<()> for Uniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, _: &mut R) -> () { () }
}
tuple_impl!{A}
tuple_impl!{A, B}
tuple_impl!{A, B, C}
tuple_impl!{A, B, C, D}
tuple_impl!{A, B, C, D, E}
tuple_impl!{A, B, C, D, E, F}
tuple_impl!{A, B, C, D, E, F, G}
tuple_impl!{A, B, C, D, E, F, G, H}
tuple_impl!{A, B, C, D, E, F, G, H, I}
tuple_impl!{A, B, C, D, E, F, G, H, I, J}
tuple_impl!{A, B, C, D, E, F, G, H, I, J, K}
tuple_impl!{A, B, C, D, E, F, G, H, I, J, K, L}

macro_rules! array_impl {
    // recursive, given at least one type parameter:
    {$n:expr, $t:ident, $($ts:ident,)*} => {
        array_impl!{($n - 1), $($ts,)*}

        impl<T> Distribution<[T; $n]> for Uniform where Uniform: Distribution<T> {
            #[inline]
            fn sample<R: Rng + ?Sized>(&self, _rng: &mut R) -> [T; $n] {
                [_rng.gen::<$t>(), $(_rng.gen::<$ts>()),*]
            }
        }
    };
    // empty case:
    {$n:expr,} => {
        impl<T> Distribution<[T; $n]> for Uniform {
            fn sample<R: Rng + ?Sized>(&self, _rng: &mut R) -> [T; $n] { [] }
        }
    };
}

array_impl!{32, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T,}

impl<T> Distribution<Option<T>> for Uniform where Uniform: Distribution<T> {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Option<T> {
        // UFCS is needed here: https://github.com/rust-lang/rust/issues/24066
        if rng.gen::<bool>() {
            Some(rng.gen())
        } else {
            None
        }
    }
}

#[cfg(test)]
mod tests {
    use {Rng, RngCore, Uniform};

    #[test]
    fn test_misc() {
        let mut rng: &mut RngCore = &mut ::test::rng(820);

        rng.sample::<char, _>(Uniform);
        rng.sample::<bool, _>(Uniform);
    }
}
<commit_msg>Generate bool using sign test<commit_after>// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// https://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The implementations of the `Uniform` distribution for other built-in types.

use core::char;

use {Rng};
use distributions::{Distribution, Uniform, Range};

impl Distribution<char> for Uniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> char {
        let range = Range::new(0u32, 0x11_0000);
        loop {
            match char::from_u32(range.sample(rng)) {
                Some(c) => return c,
                // About 0.2% of numbers in the range 0..0x110000 are invalid
                // codepoints (surrogates).
                None => {}
            }
        }
    }
}

impl Distribution<bool> for Uniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> bool {
        // We can compare against an arbitrary bit of an u32 to get a bool.
        // Because the least significant bits of a lower quality RNG can have
        // simple patterns, we compare against the most significant bit. This is
        // easiest done using a sign test.
        (rng.next_u32() as i32) < 0
    }
}

macro_rules! tuple_impl {
    // use variables to indicate the arity of the tuple
    ($($tyvar:ident),* ) => {
        // the trailing commas are for the 1 tuple
        impl< $( $tyvar ),* >
            Distribution<( $( $tyvar ),* , )>
            for Uniform
            where $( Uniform: Distribution<$tyvar> ),*
        {
            #[inline]
            fn sample<R: Rng + ?Sized>(&self, _rng: &mut R) -> ( $( $tyvar ),* , ) {
                (
                    // use the $tyvar's to get the appropriate number of
                    // repeats (they're not actually needed)
                    $(
                        _rng.gen::<$tyvar>()
                    ),*
                    ,
                )
            }
        }
    }
}

impl Distribution<()> for Uniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, _: &mut R) -> () { () }
}
tuple_impl!{A}
tuple_impl!{A, B}
tuple_impl!{A, B, C}
tuple_impl!{A, B, C, D}
tuple_impl!{A, B, C, D, E}
tuple_impl!{A, B, C, D, E, F}
tuple_impl!{A, B, C, D, E, F, G}
tuple_impl!{A, B, C, D, E, F, G, H}
tuple_impl!{A, B, C, D, E, F, G, H, I}
tuple_impl!{A, B, C, D, E, F, G, H, I, J}
tuple_impl!{A, B, C, D, E, F, G, H, I, J, K}
tuple_impl!{A, B, C, D, E, F, G, H, I, J, K, L}

macro_rules! array_impl {
    // recursive, given at least one type parameter:
    {$n:expr, $t:ident, $($ts:ident,)*} => {
        array_impl!{($n - 1), $($ts,)*}

        impl<T> Distribution<[T; $n]> for Uniform where Uniform: Distribution<T> {
            #[inline]
            fn sample<R: Rng + ?Sized>(&self, _rng: &mut R) -> [T; $n] {
                [_rng.gen::<$t>(), $(_rng.gen::<$ts>()),*]
            }
        }
    };
    // empty case:
    {$n:expr,} => {
        impl<T> Distribution<[T; $n]> for Uniform {
            fn sample<R: Rng + ?Sized>(&self, _rng: &mut R) -> [T; $n] { [] }
        }
    };
}

array_impl!{32, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T,}

impl<T> Distribution<Option<T>> for Uniform where Uniform: Distribution<T> {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Option<T> {
        // UFCS is needed here: https://github.com/rust-lang/rust/issues/24066
        if rng.gen::<bool>() {
            Some(rng.gen())
        } else {
            None
        }
    }
}

#[cfg(test)]
mod tests {
    use {Rng, RngCore, Uniform};

    #[test]
    fn test_misc() {
        let mut rng: &mut RngCore = &mut ::test::rng(820);

        rng.sample::<char, _>(Uniform);
        rng.sample::<bool, _>(Uniform);
    }
}
<|endoftext|>


<commit_before>use std::ffi::{CStr, CString};
use std::marker::PhantomData;

use tensorflow_sys as tf;

use crate::eager::{Context, ReadonlyTensor};
use crate::{AnyTensor, DataType, Result, Status, TensorType};

/// A handle to a tensor on a device.
///
/// Constructing a TensorHandle requires a reference to an execute context so that the
/// generated handle will not out live the context.
/// ```
/// # use tensorflow::{Result, Tensor};
/// use tensorflow::eager::*;
/// # fn main() -> Result<()> {
/// let opts = ContextOptions::new();
/// let ctx = Context::new(opts)?;
///
/// let t = Tensor::from(&[3i32]).freeze();
/// let h = TensorHandle::new(&ctx, &t)?;
/// let v = h.resolve::<i32>()?;
/// assert_eq!(&v[..], &[3i32]);
/// # Ok(())
/// # }
/// ```
///
/// Since TensorHandle cannot be alive beyond the lifetime of the context, the following
/// code will not compile.
/// ```compile_fail
/// # use tensorflow::{Result, Tensor};
/// use tensorflow::eager::*;
///
/// # fn main() -> Result<()> {
/// let h = {
///     let opts = ContextOptions::new();
///     let ctx = Context::new(opts)?;
///
///     let t = Tensor::from(&[3i32]).freeze();
///     let h = TensorHandle::new(&ctx, &t)?;
///     h
/// };
/// # Ok(())
/// # }
/// ```
#[derive(Debug)]
pub struct TensorHandle<'a> {
    pub(super) inner: *mut tf::TFE_TensorHandle,
    // TensorHandle should not live longer than a given context.
    ctx: PhantomData<&'a Context>,
}

impl<'a> Drop for TensorHandle<'a> {
    fn drop(&mut self) {
        unsafe {
            tf::TFE_DeleteTensorHandle(self.inner);
        }
    }
}

impl<'a> TensorHandle<'a> {
    /// Create a TensorHandle from the input Tensor
    pub fn new<T: TensorType>(
        _ctx: &'a Context,
        t: &ReadonlyTensor<T>,
    ) -> Result<TensorHandle<'a>> {
        let status = Status::new();
        let inner = unsafe { tf::TFE_NewTensorHandle(t.inner()?, status.inner) };

        if inner.is_null() {
            Err(status)
        } else {
            Ok(TensorHandle {
                inner,
                ctx: PhantomData,
            })
        }
    }

    /// Return the DataType that corresponds to this type.
    pub fn data_type(&self) -> DataType {
        unsafe { DataType::from_c(tf::TFE_TensorHandleDataType(self.inner)) }
    }

    /// Return the number of dimensions.
    ///
    /// This function will block till the operation that produces the TensorHandle has completed.
    pub fn num_dims(&self) -> Result<usize> {
        let status = Status::new();
        let num_dims = unsafe { tf::TFE_TensorHandleNumDims(self.inner, status.inner) };
        if status.is_ok() {
            // num_dims >= 0 when the status is ok, so we can safely cast it to u64.
            Ok(num_dims as usize)
        } else {
            Err(status)
        }
    }

    /// Return the number of elements
    pub fn num_elements(&self) -> Result<u64> {
        let status = Status::new();
        let num_elements = unsafe { tf::TFE_TensorHandleNumElements(self.inner, status.inner) };
        if status.is_ok() {
            // num_elements >= 0 when the status is ok, so we can safely cast it to u64.
            Ok(num_elements as u64)
        } else {
            Err(status)
        }
    }

    /// Return the number of elements for a given dim_index.
    ///
    /// This function will block till the operation that produces the TensorHandle has completed.
    pub fn dim(&self, dim_index: i32) -> Result<u64> {
        let status = Status::new();
        let dim = unsafe { tf::TFE_TensorHandleDim(self.inner, dim_index, status.inner) };
        if status.is_ok() {
            // dim >= 0 when the status is ok, so we can safely cast it to u64.
            Ok(dim as u64)
        } else {
            Err(status)
        }
    }

    /// Return the device of the operation that produced the current TensorHandle.
    ///
    /// If the TensorHandle was produced by a copy, returns the destination device of the copy.
    /// Note that the returned device name is not always the device holding the tensor handle's memory.
    /// If you want the latter, use backing_device_name.
    ///
    /// This function will block till the operation that produces the current TensorHandle has completed.
    pub fn device_name(&self) -> Result<String> {
        let status = Status::new();
        unsafe {
            let device_name = tf::TFE_TensorHandleDeviceName(self.inner, status.inner);
            if status.is_ok() {
                Ok(CStr::from_ptr(device_name).to_str()?.to_string())
            } else {
                Err(status)
            }
        }
    }

    /// Returns the name of the device in whose memory underlying the current TensorHandle resides.
    ///
    /// This function will block till the operation that produces the current TensorHandle has completed.
    pub fn backing_device_name(&self) -> Result<String> {
        let status = Status::new();
        unsafe {
            let device_name = tf::TFE_TensorHandleBackingDeviceName(self.inner, status.inner);
            if status.is_ok() {
                Ok(CStr::from_ptr(device_name).to_str()?.to_string())
            } else {
                Err(status)
            }
        }
    }

    /// Return a new TensorHandle that shares the underlying tensor with the current TensorHandle.
    pub fn copy_sharing_tensor(&self) -> Result<Self> {
        let status = Status::new();
        let inner = unsafe { tf::TFE_TensorHandleCopySharingTensor(self.inner, status.inner) };
        if status.is_ok() {
            Ok(Self {
                inner,
                ctx: self.ctx,
            })
        } else {
            Err(status)
        }
    }

    /// This function will block till the operation that produces the current TensorHandle has completed.
    /// The memory returned might alias the internal memory used by TensorFlow.
    /// Hence, callers should not mutate this memory.
    pub fn resolve<T: TensorType>(&self) -> Result<ReadonlyTensor<T>> {
        let mut status = Status::new();
        let tf_tensor = unsafe { tf::TFE_TensorHandleResolve(self.inner, status.inner) };
        if !status.is_ok() {
            return Err(status);
        }
        if self.data_type() != T::data_type() {
            let msg = format!(
                "The expected data type ({}) and underlying data type ({}) did not match.",
                T::data_type(),
                self.data_type()
            );

            status.set_lossy(crate::Code::InvalidArgument, &msg);
            return Err(status);
        }

        // Safely unwrap since data_type was checked beforehand.
        unsafe { Ok(ReadonlyTensor::from_tf_tensor(tf_tensor).unwrap()) }
    }

    /// Create a new TensorHandle with the same contents as the current TensorHandle but placed
    /// in the memory of the device name 'device_name'.
    /// If source and destination are the same device, then this creates a new handle
    /// that shares the underlying buffer. Otherwise, it currently requires at least
    /// one of the source or destination devices to be CPU (i.e., for the source or
    /// destination tensor to be placed in host memory).
    /// If async execution is enabled, the copy may be enqueued and the call will
    /// return "non-ready" TensorHandle. Else, this function returns after the copy has
    /// been done.
    pub fn copy_to_device<'b>(
        &self,
        ctx: &'b Context,
        device_name: &str,
    ) -> Result<TensorHandle<'b>> {
        let status = Status::new();

        let device_name = CString::new(device_name)?;
        unsafe {
            let inner = tf::TFE_TensorHandleCopyToDevice(
                self.inner,
                ctx.inner,
                device_name.as_ptr(),
                status.inner,
            );

            if status.is_ok() {
                Ok(TensorHandle {
                    inner,
                    ctx: PhantomData,
                })
            } else {
                Err(status)
            }
        }
    }

    /// Convert the raw TFE_TensorHandle* into a TensorHandle.
    pub(super) unsafe fn from_tensor_handle(
        _ctx: &'a Context,
        inner: *mut tf::TFE_TensorHandle,
    ) -> Self {
        Self {
            inner,
            ctx: PhantomData,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::eager::ContextOptions;
    use crate::Tensor;

    #[test]
    fn test_tensor_handle() {
        let opts = ContextOptions::new();
        let ctx = Context::new(opts).unwrap();

        let t = Tensor::new(&[2, 3])
            .with_values(&[0_i32, 1, 2, 3, 4, 5])
            .unwrap()
            .freeze();
        let h = TensorHandle::new(&ctx, &t).unwrap();

        assert_eq!(h.data_type(), DataType::Int32);
        assert_eq!(h.num_elements().unwrap(), 6);
        assert_eq!(h.num_dims().unwrap(), 2);
        assert_eq!(h.dim(0).unwrap(), 2);
        assert_eq!(h.dim(1).unwrap(), 3);
    }

    #[test]
    fn test_copy_sharing_tensor() {
        let opts = ContextOptions::new();
        let ctx = Context::new(opts).unwrap();

        let t = Tensor::new(&[2, 3])
            .with_values(&[0_i32, 1, 2, 3, 4, 5])
            .unwrap()
            .freeze();
        let h = TensorHandle::new(&ctx, &t).unwrap();
        let h_copy = h.copy_sharing_tensor().unwrap();
        let t2 = h_copy.resolve::<i32>().unwrap();

        // t and t2 may share the same memory, but it's difficuly to check
        // since the `resolve` does not guarantee that.
        assert_eq!(t, t2);
    }

    /// Following tests are disabled by default because it requires a GPU and some setup.
    ///
    /// To run this test, you need to pass the `-- --ignored` argument to cargo test.
    /// ```sh
    /// cargo test --features "eager tensorflow_gpu" -- --ignored
    /// ```
    #[cfg(feature = "tensorflow_gpu")]
    mod gpu {
        use super::*;
        use crate::eager::ContextOptions;

        #[test]
        #[ignore]
        fn test_copy_to_device() {
            let values = [0_i32, 1, 2, 3];

            let opts = ContextOptions::new();
            let ctx = Context::new(opts).unwrap();
            let devices = ctx.device_list().unwrap();
            let gpu_device = devices
                .iter()
                .find(|d| d.device_type == "GPU")
                .expect("No GPU device was found.");
            let target_device = &gpu_device.name;

            let t = Tensor::new(&[2, 2]).with_values(&values).unwrap();
            let h = TensorHandle::new(&ctx, &t).unwrap();
            let h_gpu = TensorHandle::copy_to_device(&h, &ctx, target_device).unwrap();
            assert_eq!(&h_gpu.device_name().unwrap(), target_device);
            let t2 = h_gpu.resolve::<i32>().unwrap();

            assert_eq!(&t[..], &t2[..]);
        }

        #[test]
        #[ignore]
        fn test_copy_to_device_lifetime() {
            let values = [0_i32, 1, 2, 3];

            let opts = ContextOptions::new();
            let ctx = Context::new(opts).unwrap();
            let devices = ctx.device_list().unwrap();
            let gpu_device = devices
                .iter()
                .find(|d| d.device_type == "GPU")
                .expect("No GPU device was found.");
            let target_device = &gpu_device.name;

            let h_gpu = {
                // Create a temporal Context
                let opts = ContextOptions::new();
                let ctx2 = Context::new(opts).unwrap();
                let t = Tensor::new(&[2, 2]).with_values(&values).unwrap();

                // Create a TensorHandle managed by the context `ctx2`.
                let h = TensorHandle::new(&ctx2, &t).unwrap();

                // Copy to GPU. This creates a new handle managed by the context `ctx`.
                h.copy_to_device(&ctx, target_device).unwrap()
            };
            assert_eq!(&h_gpu.device_name().unwrap(), target_device);
            let t2 = h_gpu.resolve::<i32>().unwrap();

            assert_eq!(&values[..], &t2[..]);
        }
    }
}
<commit_msg>update TensorHandle's doc<commit_after>use std::ffi::{CStr, CString};
use std::marker::PhantomData;

use tensorflow_sys as tf;

use crate::eager::{Context, ReadonlyTensor};
use crate::{AnyTensor, DataType, Result, Status, TensorType};

/// A handle to a tensor on a device.
///
/// Constructing a TensorHandle requires a reference to an execute context so that the
/// generated handle will not out live the context.
/// ```
/// # use tensorflow::{Result, Tensor};
/// use tensorflow::eager::*;
/// # fn main() -> Result<()> {
/// let opts = ContextOptions::new();
/// let ctx = Context::new(opts)?;
///
/// let t = Tensor::from(&[3i32]).freeze();
/// let h = TensorHandle::new(&ctx, &t)?;
/// let v = h.resolve::<i32>()?;
/// assert_eq!(&v[..], &[3i32]);
/// # Ok(())
/// # }
/// ```
///
/// TensorHandle manages the same buffer of the tensor. Users can destruct the Tensor
/// while leaving the TensorHandle. This is a valid use case for TensorHandle.
/// ```
/// # use tensorflow::{Result, Tensor};
/// use tensorflow::eager::*;
///
/// # fn main() -> Result<()> {
/// let opts = ContextOptions::new();
/// let ctx = Context::new(opts)?;
/// let h = {
///     let t = Tensor::from(&[3i32]).freeze();
///     TensorHandle::new(&ctx, &t)?
/// };
/// // At this point, the buffer is managed only by the handle.
/// # Ok(())
/// # }
/// ```
///
/// Since TensorHandle cannot be alive beyond the lifetime of the context, the following
/// code will not compile.
/// ```compile_fail
/// # use tensorflow::{Result, Tensor};
/// use tensorflow::eager::*;
///
/// # fn main() -> Result<()> {
/// let h = {
///     let opts = ContextOptions::new();
///     let ctx = Context::new(opts)?;
///
///     let t = Tensor::from(&[3i32]).freeze();
///     TensorHandle::new(&ctx, &t)?
/// };
/// # Ok(())
/// # }
/// ```
///
#[derive(Debug)]
pub struct TensorHandle<'a> {
    pub(super) inner: *mut tf::TFE_TensorHandle,
    // TensorHandle should not live longer than a given context.
    ctx: PhantomData<&'a Context>,
}

impl<'a> Drop for TensorHandle<'a> {
    fn drop(&mut self) {
        unsafe {
            tf::TFE_DeleteTensorHandle(self.inner);
        }
    }
}

impl<'a> TensorHandle<'a> {
    /// Create a TensorHandle from the input Tensor
    pub fn new<T: TensorType>(
        _ctx: &'a Context,
        t: &ReadonlyTensor<T>,
    ) -> Result<TensorHandle<'a>> {
        let status = Status::new();
        let inner = unsafe { tf::TFE_NewTensorHandle(t.inner()?, status.inner) };

        if inner.is_null() {
            Err(status)
        } else {
            Ok(TensorHandle {
                inner,
                ctx: PhantomData,
            })
        }
    }

    /// Return the DataType that corresponds to this type.
    pub fn data_type(&self) -> DataType {
        unsafe { DataType::from_c(tf::TFE_TensorHandleDataType(self.inner)) }
    }

    /// Return the number of dimensions.
    ///
    /// This function will block till the operation that produces the TensorHandle has completed.
    pub fn num_dims(&self) -> Result<usize> {
        let status = Status::new();
        let num_dims = unsafe { tf::TFE_TensorHandleNumDims(self.inner, status.inner) };
        if status.is_ok() {
            // num_dims >= 0 when the status is ok, so we can safely cast it to u64.
            Ok(num_dims as usize)
        } else {
            Err(status)
        }
    }

    /// Return the number of elements
    pub fn num_elements(&self) -> Result<u64> {
        let status = Status::new();
        let num_elements = unsafe { tf::TFE_TensorHandleNumElements(self.inner, status.inner) };
        if status.is_ok() {
            // num_elements >= 0 when the status is ok, so we can safely cast it to u64.
            Ok(num_elements as u64)
        } else {
            Err(status)
        }
    }

    /// Return the number of elements for a given dim_index.
    ///
    /// This function will block till the operation that produces the TensorHandle has completed.
    pub fn dim(&self, dim_index: i32) -> Result<u64> {
        let status = Status::new();
        let dim = unsafe { tf::TFE_TensorHandleDim(self.inner, dim_index, status.inner) };
        if status.is_ok() {
            // dim >= 0 when the status is ok, so we can safely cast it to u64.
            Ok(dim as u64)
        } else {
            Err(status)
        }
    }

    /// Return the device of the operation that produced the current TensorHandle.
    ///
    /// If the TensorHandle was produced by a copy, returns the destination device of the copy.
    /// Note that the returned device name is not always the device holding the tensor handle's memory.
    /// If you want the latter, use backing_device_name.
    ///
    /// This function will block till the operation that produces the current TensorHandle has completed.
    pub fn device_name(&self) -> Result<String> {
        let status = Status::new();
        unsafe {
            let device_name = tf::TFE_TensorHandleDeviceName(self.inner, status.inner);
            if status.is_ok() {
                Ok(CStr::from_ptr(device_name).to_str()?.to_string())
            } else {
                Err(status)
            }
        }
    }

    /// Returns the name of the device in whose memory underlying the current TensorHandle resides.
    ///
    /// This function will block till the operation that produces the current TensorHandle has completed.
    pub fn backing_device_name(&self) -> Result<String> {
        let status = Status::new();
        unsafe {
            let device_name = tf::TFE_TensorHandleBackingDeviceName(self.inner, status.inner);
            if status.is_ok() {
                Ok(CStr::from_ptr(device_name).to_str()?.to_string())
            } else {
                Err(status)
            }
        }
    }

    /// Return a new TensorHandle that shares the underlying tensor with the current TensorHandle.
    pub fn copy_sharing_tensor(&self) -> Result<Self> {
        let status = Status::new();
        let inner = unsafe { tf::TFE_TensorHandleCopySharingTensor(self.inner, status.inner) };
        if status.is_ok() {
            Ok(Self {
                inner,
                ctx: self.ctx,
            })
        } else {
            Err(status)
        }
    }

    /// This function will block till the operation that produces the current TensorHandle has completed.
    /// The memory returned might alias the internal memory used by TensorFlow.
    /// Hence, callers should not mutate this memory.
    pub fn resolve<T: TensorType>(&self) -> Result<ReadonlyTensor<T>> {
        let mut status = Status::new();
        let tf_tensor = unsafe { tf::TFE_TensorHandleResolve(self.inner, status.inner) };
        if !status.is_ok() {
            return Err(status);
        }
        if self.data_type() != T::data_type() {
            let msg = format!(
                "The expected data type ({}) and underlying data type ({}) did not match.",
                T::data_type(),
                self.data_type()
            );

            status.set_lossy(crate::Code::InvalidArgument, &msg);
            return Err(status);
        }

        // Safely unwrap since data_type was checked beforehand.
        unsafe { Ok(ReadonlyTensor::from_tf_tensor(tf_tensor).unwrap()) }
    }

    /// Create a new TensorHandle with the same contents as the current TensorHandle but placed
    /// in the memory of the device name 'device_name'.
    /// If source and destination are the same device, then this creates a new handle
    /// that shares the underlying buffer. Otherwise, it currently requires at least
    /// one of the source or destination devices to be CPU (i.e., for the source or
    /// destination tensor to be placed in host memory).
    /// If async execution is enabled, the copy may be enqueued and the call will
    /// return "non-ready" TensorHandle. Else, this function returns after the copy has
    /// been done.
    pub fn copy_to_device<'b>(
        &self,
        ctx: &'b Context,
        device_name: &str,
    ) -> Result<TensorHandle<'b>> {
        let status = Status::new();

        let device_name = CString::new(device_name)?;
        unsafe {
            let inner = tf::TFE_TensorHandleCopyToDevice(
                self.inner,
                ctx.inner,
                device_name.as_ptr(),
                status.inner,
            );

            if status.is_ok() {
                Ok(TensorHandle {
                    inner,
                    ctx: PhantomData,
                })
            } else {
                Err(status)
            }
        }
    }

    /// Convert the raw TFE_TensorHandle* into a TensorHandle.
    pub(super) unsafe fn from_tensor_handle(
        _ctx: &'a Context,
        inner: *mut tf::TFE_TensorHandle,
    ) -> Self {
        Self {
            inner,
            ctx: PhantomData,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::eager::ContextOptions;
    use crate::Tensor;

    #[test]
    fn test_tensor_handle() {
        let opts = ContextOptions::new();
        let ctx = Context::new(opts).unwrap();

        let t = Tensor::new(&[2, 3])
            .with_values(&[0_i32, 1, 2, 3, 4, 5])
            .unwrap()
            .freeze();
        let h = TensorHandle::new(&ctx, &t).unwrap();

        assert_eq!(h.data_type(), DataType::Int32);
        assert_eq!(h.num_elements().unwrap(), 6);
        assert_eq!(h.num_dims().unwrap(), 2);
        assert_eq!(h.dim(0).unwrap(), 2);
        assert_eq!(h.dim(1).unwrap(), 3);
    }

    #[test]
    fn test_copy_sharing_tensor() {
        let opts = ContextOptions::new();
        let ctx = Context::new(opts).unwrap();

        let t = Tensor::new(&[2, 3])
            .with_values(&[0_i32, 1, 2, 3, 4, 5])
            .unwrap()
            .freeze();
        let h = TensorHandle::new(&ctx, &t).unwrap();
        let h_copy = h.copy_sharing_tensor().unwrap();
        let t2 = h_copy.resolve::<i32>().unwrap();

        // t and t2 may share the same memory, but it's difficuly to check
        // since the `resolve` does not guarantee that.
        assert_eq!(t, t2);
    }

    /// Following tests are disabled by default because it requires a GPU and some setup.
    ///
    /// To run this test, you need to pass the `-- --ignored` argument to cargo test.
    /// ```sh
    /// cargo test --features "eager tensorflow_gpu" -- --ignored
    /// ```
    #[cfg(feature = "tensorflow_gpu")]
    mod gpu {
        use super::*;
        use crate::eager::ContextOptions;

        #[test]
        #[ignore]
        fn test_copy_to_device() {
            let values = [0_i32, 1, 2, 3];

            let opts = ContextOptions::new();
            let ctx = Context::new(opts).unwrap();
            let devices = ctx.device_list().unwrap();
            let gpu_device = devices
                .iter()
                .find(|d| d.device_type == "GPU")
                .expect("No GPU device was found.");
            let target_device = &gpu_device.name;

            let t = Tensor::new(&[2, 2]).with_values(&values).unwrap();
            let h = TensorHandle::new(&ctx, &t).unwrap();
            let h_gpu = TensorHandle::copy_to_device(&h, &ctx, target_device).unwrap();
            assert_eq!(&h_gpu.device_name().unwrap(), target_device);
            let t2 = h_gpu.resolve::<i32>().unwrap();

            assert_eq!(&t[..], &t2[..]);
        }

        #[test]
        #[ignore]
        fn test_copy_to_device_lifetime() {
            let values = [0_i32, 1, 2, 3];

            let opts = ContextOptions::new();
            let ctx = Context::new(opts).unwrap();
            let devices = ctx.device_list().unwrap();
            let gpu_device = devices
                .iter()
                .find(|d| d.device_type == "GPU")
                .expect("No GPU device was found.");
            let target_device = &gpu_device.name;

            let h_gpu = {
                // Create a temporal Context
                let opts = ContextOptions::new();
                let ctx2 = Context::new(opts).unwrap();
                let t = Tensor::new(&[2, 2]).with_values(&values).unwrap();

                // Create a TensorHandle managed by the context `ctx2`.
                let h = TensorHandle::new(&ctx2, &t).unwrap();

                // Copy to GPU. This creates a new handle managed by the context `ctx`.
                h.copy_to_device(&ctx, target_device).unwrap()
            };
            assert_eq!(&h_gpu.device_name().unwrap(), target_device);
            let t2 = h_gpu.resolve::<i32>().unwrap();

            assert_eq!(&values[..], &t2[..]);
        }
    }
}
<|endoftext|>


<commit_before>use std::mem;
use std::ptr;

use runtime::{Class, Method, Object, Sel, self};
use {encode, Encode, EncodePtr};

/*
 The Sized bound on Message is unfortunate; ideally, objc objects would not be
 treated as Sized. However, rust won't allow casting a dynamically-sized type
 pointer to an Object pointer, because dynamically-sized types can have fat
 pointers (two words) instead of real pointers.
 */
/// Types that may be sent Objective-C messages.
/// For example: objects, classes, and blocks.
pub unsafe trait Message: Sized + EncodePtr { }

unsafe impl Message for Object { }

unsafe impl Message for Class { }

/// A trait for converting to a pointer to a type that may be sent Objective-C
/// messages.
pub trait ToMessage : Encode {
    type Target: Message;

    fn as_ptr(&self) -> *mut Self::Target;

    fn is_nil(&self) -> bool {
        self.as_ptr().is_null()
    }
}

impl<T> ToMessage for *const T where T: Message {
    type Target = T;

    fn as_ptr(&self) -> *mut T {
        *self as *mut T
    }
}

impl<T> ToMessage for *mut T where T: Message {
    type Target = T;

    fn as_ptr(&self) -> *mut T {
        *self
    }
}

impl<'a, T> ToMessage for &'a T where T: Message {
    type Target = T;

    fn as_ptr(&self) -> *mut T {
        *self as *const T as *mut T
    }
}

impl<'a, T> ToMessage for &'a mut T where T: Message {
    type Target = T;

    fn as_ptr(&self) -> *mut T {
        *self
    }
}

impl<'a, T> ToMessage for Option<&'a T> where T: Message {
    type Target = T;

    fn as_ptr(&self) -> *mut T {
        match *self {
            None => ptr::null_mut(),
            Some(ref obj) => obj.as_ptr(),
        }
    }
}

impl<'a, T> ToMessage for Option<&'a mut T> where T: Message {
    type Target = T;

    fn as_ptr(&self) -> *mut T {
        match *self {
            None => ptr::null_mut(),
            Some(ref obj) => obj.as_ptr(),
        }
    }
}

fn msg_send_fn<R>() -> unsafe extern fn(*mut Object, Sel, ...) -> R {
    unsafe {
        mem::transmute(runtime::objc_msgSend)
    }
}

pub trait MessageArguments {
    unsafe fn send<T, R>(self, obj: &T, sel: Sel) -> R where T: ToMessage;

    fn verify(&self, method: &Method) -> Result<(), String>;
}

macro_rules! message_args_impl {
    ($($a:ident : $t:ident),*) => (
        impl<$($t: Encode),*> MessageArguments for ($($t,)*) {
            unsafe fn send<T, R>(self, obj: &T, sel: Sel) -> R where T: ToMessage {
                let msg_send_fn = msg_send_fn::<R>();
                let msg_send_fn: unsafe extern fn(*mut Object, Sel $(, $t)*) -> R =
                    mem::transmute(msg_send_fn);
                let obj_ptr = obj.as_ptr() as *mut Object;
                let ($($a,)*) = self;
                msg_send_fn(obj_ptr, sel $(, $a)*)
            }

            fn verify(&self, method: &Method) -> Result<(), String> {
                let types = [$(encode::<$t>()),*];
                // So rust can infer the type of the empty array:
                let _: &[&str] = &types;

                let count = 2 + types.len();
                let expected_count = method.arguments_count();
                if count != expected_count {
                    return Err(format!("Method {:?} accepts {} arguments, but {} were given",
                        method.name(), expected_count, count));
                }

                let expected_types = range(2, expected_count).map(
                    |i| method.argument_type(i));
                for (&arg, expected) in types.iter().zip(expected_types) {
                    let expected = match expected {
                        Some(s) => s,
                        None => return Err(format!("Method {:?} doesn't expect argument with type code {}",
                            method.name(), arg)),
                    };
                    if arg != expected.as_slice() {
                        return Err(format!("Method {:?} expected argument with type code {} but was given {}",
                            method.name(), expected.as_slice(), arg));
                    }
                }

                Ok(())
            }
        }
    );
}

message_args_impl!();
message_args_impl!(a: A);
message_args_impl!(a: A, b: B);
message_args_impl!(a: A, b: B, c: C);
message_args_impl!(a: A, b: B, c: C, d: D);
message_args_impl!(a: A, b: B, c: C, d: D, e: E);
message_args_impl!(a: A, b: B, c: C, d: D, e: E, f: F);
message_args_impl!(a: A, b: B, c: C, d: D, e: E, f: F, g: G);
message_args_impl!(a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H);
message_args_impl!(a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I);
message_args_impl!(a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J);
message_args_impl!(a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J, k: K);
message_args_impl!(a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J, k: K, l: L);

fn verify_message_signature<T, A, R>(obj: Option<&T>, sel: Sel, args: &A) ->
        Result<(), String> where T: Message, A: MessageArguments, R: Encode {
    let obj = match obj {
        Some(obj) => obj,
        None => return Err(format!("Messaging {:?} to nil", sel)),
    };
    let cls = unsafe {
        let obj = &*(obj as *const _ as *const Object);
        obj.class()
    };
    let method = match cls.instance_method(sel) {
        Some(method) => method,
        None => return Err(format!("Method {:?} not found on class {:?}",
            sel, cls)),
    };

    let ret = encode::<R>();
    let expected_ret = method.return_type();
    // Allow encoding "oneway void" (Vv) as "void" (v)
    if expected_ret.as_slice() != ret && !(expected_ret.as_slice() == "Vv" && ret == "v") {
        return Err(format!("Return type code {} does not match expected {} for method {:?} on class {:?}",
            ret, expected_ret.as_slice(), sel, cls));
    }

    // I don't think either of these can happen, but just to be safe...
    let accepts_self_arg = match method.argument_type(0) {
        Some(s) => s.as_slice() == "@",
        None => false,
    };
    if !accepts_self_arg {
        return Err(format!("Method {:?} of class {:?} doesn't accept an argument for self",
            sel, cls));
    }
    let accepts_cmd_arg = match method.argument_type(1) {
        Some(s) => s.as_slice() == ":",
        None => false,
    };
    if !accepts_cmd_arg {
        return Err(format!("Method {:?} of class {:?} doesn't accept an argument for the selector",
            sel, cls));
    }

    args.verify(method)
}

pub unsafe fn send_message<T, A, R>(obj: &T, sel: Sel, args: A) -> R
        where T: ToMessage, A: MessageArguments {
    args.send(obj, sel)
}

pub unsafe fn send_message_verified<T, A, R>(obj: &T, sel: Sel, args: A) ->
        Result<R, String> where T: ToMessage, A: MessageArguments, R: Encode {
    let obj_ref = obj.as_ptr().as_ref();
    verify_message_signature::<_, _, R>(obj_ref, sel, &args).and_then(
        move |()| Ok(args.send(obj, sel)))
}
<commit_msg>Extracted argument verification into a function.<commit_after>use std::mem;
use std::ptr;

use runtime::{Class, Method, Object, Sel, self};
use {encode, Encode, EncodePtr};

/*
 The Sized bound on Message is unfortunate; ideally, objc objects would not be
 treated as Sized. However, rust won't allow casting a dynamically-sized type
 pointer to an Object pointer, because dynamically-sized types can have fat
 pointers (two words) instead of real pointers.
 */
/// Types that may be sent Objective-C messages.
/// For example: objects, classes, and blocks.
pub unsafe trait Message: Sized + EncodePtr { }

unsafe impl Message for Object { }

unsafe impl Message for Class { }

/// A trait for converting to a pointer to a type that may be sent Objective-C
/// messages.
pub trait ToMessage : Encode {
    type Target: Message;

    fn as_ptr(&self) -> *mut Self::Target;

    fn is_nil(&self) -> bool {
        self.as_ptr().is_null()
    }
}

impl<T> ToMessage for *const T where T: Message {
    type Target = T;

    fn as_ptr(&self) -> *mut T {
        *self as *mut T
    }
}

impl<T> ToMessage for *mut T where T: Message {
    type Target = T;

    fn as_ptr(&self) -> *mut T {
        *self
    }
}

impl<'a, T> ToMessage for &'a T where T: Message {
    type Target = T;

    fn as_ptr(&self) -> *mut T {
        *self as *const T as *mut T
    }
}

impl<'a, T> ToMessage for &'a mut T where T: Message {
    type Target = T;

    fn as_ptr(&self) -> *mut T {
        *self
    }
}

impl<'a, T> ToMessage for Option<&'a T> where T: Message {
    type Target = T;

    fn as_ptr(&self) -> *mut T {
        match *self {
            None => ptr::null_mut(),
            Some(ref obj) => obj.as_ptr(),
        }
    }
}

impl<'a, T> ToMessage for Option<&'a mut T> where T: Message {
    type Target = T;

    fn as_ptr(&self) -> *mut T {
        match *self {
            None => ptr::null_mut(),
            Some(ref obj) => obj.as_ptr(),
        }
    }
}

fn msg_send_fn<R>() -> unsafe extern fn(*mut Object, Sel, ...) -> R {
    unsafe {
        mem::transmute(runtime::objc_msgSend)
    }
}

pub trait MessageArguments {
    unsafe fn send<T, R>(self, obj: &T, sel: Sel) -> R where T: ToMessage;

    fn verify(&self, method: &Method) -> Result<(), String>;
}

macro_rules! message_args_impl {
    ($($a:ident : $t:ident),*) => (
        impl<$($t: Encode),*> MessageArguments for ($($t,)*) {
            unsafe fn send<T, R>(self, obj: &T, sel: Sel) -> R where T: ToMessage {
                let msg_send_fn = msg_send_fn::<R>();
                let msg_send_fn: unsafe extern fn(*mut Object, Sel $(, $t)*) -> R =
                    mem::transmute(msg_send_fn);
                let obj_ptr = obj.as_ptr() as *mut Object;
                let ($($a,)*) = self;
                msg_send_fn(obj_ptr, sel $(, $a)*)
            }

            fn verify(&self, method: &Method) -> Result<(), String> {
                let types = [$(encode::<$t>()),*];
                verify_message_arguments(&types, method)
            }
        }
    );
}

message_args_impl!();
message_args_impl!(a: A);
message_args_impl!(a: A, b: B);
message_args_impl!(a: A, b: B, c: C);
message_args_impl!(a: A, b: B, c: C, d: D);
message_args_impl!(a: A, b: B, c: C, d: D, e: E);
message_args_impl!(a: A, b: B, c: C, d: D, e: E, f: F);
message_args_impl!(a: A, b: B, c: C, d: D, e: E, f: F, g: G);
message_args_impl!(a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H);
message_args_impl!(a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I);
message_args_impl!(a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J);
message_args_impl!(a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J, k: K);
message_args_impl!(a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J, k: K, l: L);

fn verify_message_arguments(types: &[&str], method: &Method) -> Result<(), String> {
    let count = 2 + types.len();
    let expected_count = method.arguments_count();
    if count != expected_count {
        return Err(format!("Method {:?} accepts {} arguments, but {} were given",
            method.name(), expected_count, count));
    }

    let expected_types = range(2, expected_count).map(|i| method.argument_type(i));
    for (&arg, expected) in types.iter().zip(expected_types) {
        let expected = match expected {
            Some(s) => s,
            None => return Err(format!("Method {:?} doesn't expect argument with type code {}",
                method.name(), arg)),
        };
        if arg != expected.as_slice() {
            return Err(format!("Method {:?} expected argument with type code {} but was given {}",
                method.name(), expected.as_slice(), arg));
        }
    }

    Ok(())
}

fn verify_message_signature<T, A, R>(obj: Option<&T>, sel: Sel, args: &A) ->
        Result<(), String> where T: Message, A: MessageArguments, R: Encode {
    let obj = match obj {
        Some(obj) => obj,
        None => return Err(format!("Messaging {:?} to nil", sel)),
    };
    let cls = unsafe {
        let obj = &*(obj as *const _ as *const Object);
        obj.class()
    };
    let method = match cls.instance_method(sel) {
        Some(method) => method,
        None => return Err(format!("Method {:?} not found on class {:?}",
            sel, cls)),
    };

    let ret = encode::<R>();
    let expected_ret = method.return_type();
    // Allow encoding "oneway void" (Vv) as "void" (v)
    if expected_ret.as_slice() != ret && !(expected_ret.as_slice() == "Vv" && ret == "v") {
        return Err(format!("Return type code {} does not match expected {} for method {:?} on class {:?}",
            ret, expected_ret.as_slice(), sel, cls));
    }

    // I don't think either of these can happen, but just to be safe...
    let accepts_self_arg = match method.argument_type(0) {
        Some(s) => s.as_slice() == "@",
        None => false,
    };
    if !accepts_self_arg {
        return Err(format!("Method {:?} of class {:?} doesn't accept an argument for self",
            sel, cls));
    }
    let accepts_cmd_arg = match method.argument_type(1) {
        Some(s) => s.as_slice() == ":",
        None => false,
    };
    if !accepts_cmd_arg {
        return Err(format!("Method {:?} of class {:?} doesn't accept an argument for the selector",
            sel, cls));
    }

    args.verify(method)
}

pub unsafe fn send_message<T, A, R>(obj: &T, sel: Sel, args: A) -> R
        where T: ToMessage, A: MessageArguments {
    args.send(obj, sel)
}

pub unsafe fn send_message_verified<T, A, R>(obj: &T, sel: Sel, args: A) ->
        Result<R, String> where T: ToMessage, A: MessageArguments, R: Encode {
    let obj_ref = obj.as_ptr().as_ref();
    verify_message_signature::<_, _, R>(obj_ref, sel, &args).and_then(
        move |()| Ok(args.send(obj, sel)))
}
<|endoftext|>


<commit_before>//! The master module provides the Master struct which instructs Workers which
//! portions of the image they should render and collects their results to combine
//! into the final image.

use std::path::PathBuf;
use std::io::prelude::*;
use std::collections::HashMap;
use std::net::ToSocketAddrs;
use std::iter;

use bincode::SizeLimit;
use bincode::rustc_serialize::{encode, decode};
use image;
use mio::tcp::{TcpStream, Shutdown};
use mio::*;

use film::Image;
use exec::Config;
use exec::distrib::{worker, Instructions, Frame};
use sampler::BlockQueue;

/// Stores distributed rendering status. The frame is either InProgress and contains
/// partially rendered results from the workers who've reported the frame or is Completed
/// and has been saved out to disk.
#[derive(Debug)]
enum DistributedFrame {
    InProgress {
        // Which frame number this is
        frame: usize,
        // The number of workers who have reported results for this
        // frame so far
        num_reporting: usize,
        render: Image,
    },
    Completed,
}

impl DistributedFrame {
    pub fn start(frame_num: usize, img_dim: (usize, usize)) -> DistributedFrame {
        DistributedFrame::InProgress { frame: frame_num, num_reporting: 0, render: Image::new(img_dim) }
    }
}

/// Buffer for collecting results from a worker asynchronously. The buffer is filled
/// as we get readable events from the workers until it reaches the expected size.
/// After this the Frame is decoded and accumulated in the appropriate DistributedFrame
#[derive(Clone, Debug)]
struct WorkerBuffer {
    pub buf: Vec<u8>,
    pub expected_size: usize,
    pub currently_read: usize,
}

impl WorkerBuffer {
    pub fn new() -> WorkerBuffer {
        WorkerBuffer { buf: Vec::new(), expected_size: 8, currently_read: 0 }
    }
}

/// The Master organizes the set of Worker processes and instructions them what parts
/// of the scene to render. As workers report results the master collects them and
/// saves out the PNG once all workers have reported the frame.
pub struct Master {
    /// Hostnames of the workers to send work too
    workers: Vec<String>,
    connections: Vec<TcpStream>,
    /// Temporary buffers to store worker results in as they're
    /// read in over TCP
    worker_buffers: Vec<WorkerBuffer>,
    config: Config,
    /// List of the frames we're collecting or have completed
    frames: HashMap<usize, DistributedFrame>,
    img_dim: (usize, usize),
    /// Number of 8x8 blocks we're assigning per worker
    blocks_per_worker: usize,
    /// Remainder of blocks that will be tacked on to the last
    /// worker's assignment
    blocks_remainder: usize,
}

impl Master {
    /// Create a new master that will contact the worker nodes passed and
    /// send instructions on what parts of the scene to start rendering
    pub fn start_workers(workers: Vec<String>, config: Config, img_dim: (usize, usize))
                         -> (Master, EventLoop<Master>) {
        // Figure out how many blocks we have for this image and assign them to our workers
        let queue = BlockQueue::new((img_dim.0 as u32, img_dim.1 as u32), (8, 8), (0, 0));
        let blocks_per_worker = queue.len() / workers.len();
        let blocks_remainder = queue.len() % workers.len();

        let mut event_loop = EventLoop::<Master>::new().unwrap();
        let mut connections = Vec::new();

        // Connect to each worker and send instructions on what to render
        // I guess this should also be queue'd up and the writes actually performed asynchronously?
        // Otherwise we need to loop on write_all while waiting for the connection to open
        for (i, host) in workers.iter().enumerate() {
            let addr = (&host[..], worker::PORT).to_socket_addrs().unwrap().next().unwrap();
            match TcpStream::connect(&addr) {
                Ok(stream) => {
                    // Each worker is identified in the event loop by their index in the vec
                    match event_loop.register(&stream, Token(i), EventSet::all(), PollOpt::level()){
                        Err(e) => println!("Error registering stream from {}: {}", host, e),
                        Ok(_) => {},
                    }
                    connections.push(stream);
                },
                Err(e) => println!("Failed to contact worker {}: {:?}", host, e),
            }
        }
        let worker_buffers: Vec<_> = iter::repeat(WorkerBuffer::new()).take(workers.len()).collect();
        let master = Master { workers: workers, connections: connections,
                              worker_buffers: worker_buffers, config: config,
                              frames: HashMap::new(),
                              img_dim: img_dim,
                              blocks_per_worker: blocks_per_worker,
                              blocks_remainder: blocks_remainder };
        (master, event_loop)
    }
    /// Read a result frame from a worker and save it into the list of frames we're collecting from
    /// all workers. Will save out the final render if all workers have reported results for this
    /// frame.
    fn save_results(&mut self, frame: Frame) {
        let frame_num = frame.frame as usize;
        let img_dim = self.img_dim;
        // Find the frame being reported and create it if we haven't received parts of this frame yet
        let mut df = self.frames.entry(frame_num).or_insert_with(
                        || DistributedFrame::start(frame_num, img_dim));

        let mut finished = false;
        match df {
            &mut DistributedFrame::InProgress { frame: _, ref mut num_reporting, ref mut render } => {
                // Collect results from the worker and see if we've finished the frame and can save
                // it out
                render.add_blocks(frame.block_size, &frame.blocks, &frame.pixels);
                *num_reporting += 1;
                if *num_reporting == self.workers.len() {
                    let out_file = match self.config.out_path.extension() {
                        Some(_) => self.config.out_path.clone(),
                        None => self.config.out_path.join(
                            PathBuf::from(format!("frame{:05}.png", frame_num))),
                    };
                    let img = render.get_srgb8();
                    let dim = render.dimensions();
                    match image::save_buffer(&out_file.as_path(), &img[..], dim.0 as u32,
                    dim.1 as u32, image::RGB(8)) {
                        Ok(_) => {},
                        Err(e) => println!("Error saving image, {}", e),
                    };
                    println!("Frame {}: rendered to '{}'\n--------------------",
                             frame_num, out_file.display());
                    finished = true;
                }
            },
            &mut DistributedFrame::Completed => println!("Worker reporting on completed frame {}?", frame_num),
        }
        // This is a bit awkward, since we borrow df in the match we can't mark it finished in there
        if finished {
            *df = DistributedFrame::Completed;
        }
    }
    /// Read results from a worker and accumulate this data in its worker buffer. Returns true if
    /// we've read the data being sent and can decode the buffer
    fn read_worker_buffer(&mut self, worker: usize) -> bool {
        let mut buf = &mut self.worker_buffers[worker];
        // If we haven't read the size of data being sent, read that now
        if buf.currently_read < 8 {
            // First 8 bytes are a u64 specifying the number of bytes being sent
            buf.buf.extend(iter::repeat(0u8).take(8));
            match self.connections[worker].read(&mut buf.buf[buf.currently_read..]) {
                Ok(n) => buf.currently_read += n,
                Err(e) => println!("Error reading results from worker {}: {}", self.workers[worker], e),
            }
            if buf.currently_read == buf.expected_size {
                // How many bytes we expect to get from the worker for a frame
                buf.expected_size = decode(&buf.buf[..]).unwrap();
                // Extend the Vec so we've got enough room for the remaning bytes, minus the 8 for the
                // encoded size header
                buf.buf.extend(iter::repeat(0u8).take(buf.expected_size - 8));
            }
        }
        // If we've finished reading the size header we can now start reading the frame data
        if buf.currently_read >= 8 {
            match self.connections[worker].read(&mut buf.buf[buf.currently_read..]) {
                Ok(n) => buf.currently_read += n,
                Err(e) => println!("Error reading results from worker {}: {}", self.workers[worker], e),
            }
        }
        buf.currently_read == buf.expected_size
    }
}

impl Handler for Master {
    type Timeout = ();
    type Message = ();

    fn ready(&mut self, event_loop: &mut EventLoop<Master>, token: Token, event: EventSet) {
        let worker = token.as_usize();
        if event.is_error() {
            println!("Error connecting too {}", self.workers[worker]);
            match self.connections[worker].shutdown(Shutdown::Both) {
                Err(e) => println!("Error shutting down worker {}: {}", worker, e),
                Ok(_) => {},
            }
            // Remove the connection from the event loop
            match event_loop.deregister(&self.connections[worker]) {
                Err(e) => println!("Error deregistering worker {}: {}", worker, e),
                Ok(_) => {},
            }
        }
        // If the worker has terminated, shutdown the read end of the connection
        if event.is_hup() {
            match self.connections[worker].shutdown(Shutdown::Both) {
                Err(e) => println!("Error shutting down worker {}: {}", worker, e),
                Ok(_) => {},
            }
            // Remove the connection from the event loop
            match event_loop.deregister(&self.connections[worker]) {
                Err(e) => println!("Error deregistering worker {}: {}", worker, e),
                Ok(_) => {},
            }
        }
        // A worker is ready to receive instructions from us
        if event.is_writable() {
            let b_start = worker * self.blocks_per_worker;
            let b_count =
                if worker == self.workers.len() - 1 {
                    self.blocks_per_worker + self.blocks_remainder
                } else {
                    self.blocks_per_worker
                };
            let instr = Instructions::new(&self.config.scene_file,
                                          (self.config.frame_info.start, self.config.frame_info.end),
                                          b_start, b_count);
            // Encode and send our instructions to the worker
            let bytes = encode(&instr, SizeLimit::Infinite).unwrap();
            match self.connections[worker].write_all(&bytes[..]) {
                Err(e) => println!("Failed to send instructions to {}: {:?}", self.workers[worker], e),
                Ok(_) => {},
            }
            // Register that we no longer care about writable events on this connection
            event_loop.reregister(&self.connections[worker], token,
                                  EventSet::readable() | EventSet::error() | EventSet::hup(),
                                  PollOpt::level()).expect("Re-registering failed");
        }
        // Some results are available from a worker
        if event.is_readable() {
            // Read results from the worker, if we've accumulated all the data being sent
            // decode and accumulate the frame
            if self.read_worker_buffer(worker) {
                let frame: Frame = decode(&self.worker_buffers[worker].buf[..]).unwrap();
                self.save_results(frame);
                // Clean up the worker buffer for the next frame
                self.worker_buffers[worker].buf.clear();
                self.worker_buffers[worker].expected_size = 8;
                self.worker_buffers[worker].currently_read = 0;
            }
        }
        // After getting results from the worker we check if we've completed all our frames
        // and exit if so
        let all_complete = self.frames.values().fold(true,
                                |all, v| {
                                    match v {
                                        &DistributedFrame::Completed => true && all,
                                        _ => false,
                                    }
                                });
        // The frame start/end range is inclusive, so we must add 1 here
        let num_frames = self.config.frame_info.end - self.config.frame_info.start + 1;
        if self.frames.len() == num_frames && all_complete {
            event_loop.shutdown();
        }
    }
}

<commit_msg>Correct outdated comment<commit_after>//! The master module provides the Master struct which instructs Workers which
//! portions of the image they should render and collects their results to combine
//! into the final image.

use std::path::PathBuf;
use std::io::prelude::*;
use std::collections::HashMap;
use std::net::ToSocketAddrs;
use std::iter;

use bincode::SizeLimit;
use bincode::rustc_serialize::{encode, decode};
use image;
use mio::tcp::{TcpStream, Shutdown};
use mio::*;

use film::Image;
use exec::Config;
use exec::distrib::{worker, Instructions, Frame};
use sampler::BlockQueue;

/// Stores distributed rendering status. The frame is either InProgress and contains
/// partially rendered results from the workers who've reported the frame or is Completed
/// and has been saved out to disk.
#[derive(Debug)]
enum DistributedFrame {
    InProgress {
        // Which frame number this is
        frame: usize,
        // The number of workers who have reported results for this
        // frame so far
        num_reporting: usize,
        render: Image,
    },
    Completed,
}

impl DistributedFrame {
    pub fn start(frame_num: usize, img_dim: (usize, usize)) -> DistributedFrame {
        DistributedFrame::InProgress { frame: frame_num, num_reporting: 0, render: Image::new(img_dim) }
    }
}

/// Buffer for collecting results from a worker asynchronously. The buffer is filled
/// as we get readable events from the workers until it reaches the expected size.
/// After this the Frame is decoded and accumulated in the appropriate DistributedFrame
#[derive(Clone, Debug)]
struct WorkerBuffer {
    pub buf: Vec<u8>,
    pub expected_size: usize,
    pub currently_read: usize,
}

impl WorkerBuffer {
    pub fn new() -> WorkerBuffer {
        WorkerBuffer { buf: Vec::new(), expected_size: 8, currently_read: 0 }
    }
}

/// The Master organizes the set of Worker processes and instructions them what parts
/// of the scene to render. As workers report results the master collects them and
/// saves out the PNG once all workers have reported the frame.
pub struct Master {
    /// Hostnames of the workers to send work too
    workers: Vec<String>,
    connections: Vec<TcpStream>,
    /// Temporary buffers to store worker results in as they're
    /// read in over TCP
    worker_buffers: Vec<WorkerBuffer>,
    config: Config,
    /// List of the frames we're collecting or have completed
    frames: HashMap<usize, DistributedFrame>,
    img_dim: (usize, usize),
    /// Number of 8x8 blocks we're assigning per worker
    blocks_per_worker: usize,
    /// Remainder of blocks that will be tacked on to the last
    /// worker's assignment
    blocks_remainder: usize,
}

impl Master {
    /// Create a new master that will contact the worker nodes passed and
    /// send instructions on what parts of the scene to start rendering
    pub fn start_workers(workers: Vec<String>, config: Config, img_dim: (usize, usize))
                         -> (Master, EventLoop<Master>) {
        // Figure out how many blocks we have for this image and assign them to our workers
        let queue = BlockQueue::new((img_dim.0 as u32, img_dim.1 as u32), (8, 8), (0, 0));
        let blocks_per_worker = queue.len() / workers.len();
        let blocks_remainder = queue.len() % workers.len();

        let mut event_loop = EventLoop::<Master>::new().unwrap();
        let mut connections = Vec::new();

        // Connect to each worker and add them to the event loop
        for (i, host) in workers.iter().enumerate() {
            let addr = (&host[..], worker::PORT).to_socket_addrs().unwrap().next().unwrap();
            match TcpStream::connect(&addr) {
                Ok(stream) => {
                    // Each worker is identified in the event loop by their index in the vec
                    match event_loop.register(&stream, Token(i), EventSet::all(), PollOpt::level()){
                        Err(e) => println!("Error registering stream from {}: {}", host, e),
                        Ok(_) => {},
                    }
                    connections.push(stream);
                },
                Err(e) => println!("Failed to contact worker {}: {:?}", host, e),
            }
        }
        let worker_buffers: Vec<_> = iter::repeat(WorkerBuffer::new()).take(workers.len()).collect();
        let master = Master { workers: workers, connections: connections,
                              worker_buffers: worker_buffers, config: config,
                              frames: HashMap::new(),
                              img_dim: img_dim,
                              blocks_per_worker: blocks_per_worker,
                              blocks_remainder: blocks_remainder };
        (master, event_loop)
    }
    /// Read a result frame from a worker and save it into the list of frames we're collecting from
    /// all workers. Will save out the final render if all workers have reported results for this
    /// frame.
    fn save_results(&mut self, frame: Frame) {
        let frame_num = frame.frame as usize;
        let img_dim = self.img_dim;
        // Find the frame being reported and create it if we haven't received parts of this frame yet
        let mut df = self.frames.entry(frame_num).or_insert_with(
                        || DistributedFrame::start(frame_num, img_dim));

        let mut finished = false;
        match df {
            &mut DistributedFrame::InProgress { frame: _, ref mut num_reporting, ref mut render } => {
                // Collect results from the worker and see if we've finished the frame and can save
                // it out
                render.add_blocks(frame.block_size, &frame.blocks, &frame.pixels);
                *num_reporting += 1;
                if *num_reporting == self.workers.len() {
                    let out_file = match self.config.out_path.extension() {
                        Some(_) => self.config.out_path.clone(),
                        None => self.config.out_path.join(
                            PathBuf::from(format!("frame{:05}.png", frame_num))),
                    };
                    let img = render.get_srgb8();
                    let dim = render.dimensions();
                    match image::save_buffer(&out_file.as_path(), &img[..], dim.0 as u32,
                    dim.1 as u32, image::RGB(8)) {
                        Ok(_) => {},
                        Err(e) => println!("Error saving image, {}", e),
                    };
                    println!("Frame {}: rendered to '{}'\n--------------------",
                             frame_num, out_file.display());
                    finished = true;
                }
            },
            &mut DistributedFrame::Completed => println!("Worker reporting on completed frame {}?", frame_num),
        }
        // This is a bit awkward, since we borrow df in the match we can't mark it finished in there
        if finished {
            *df = DistributedFrame::Completed;
        }
    }
    /// Read results from a worker and accumulate this data in its worker buffer. Returns true if
    /// we've read the data being sent and can decode the buffer
    fn read_worker_buffer(&mut self, worker: usize) -> bool {
        let mut buf = &mut self.worker_buffers[worker];
        // If we haven't read the size of data being sent, read that now
        if buf.currently_read < 8 {
            // First 8 bytes are a u64 specifying the number of bytes being sent
            buf.buf.extend(iter::repeat(0u8).take(8));
            match self.connections[worker].read(&mut buf.buf[buf.currently_read..]) {
                Ok(n) => buf.currently_read += n,
                Err(e) => println!("Error reading results from worker {}: {}", self.workers[worker], e),
            }
            if buf.currently_read == buf.expected_size {
                // How many bytes we expect to get from the worker for a frame
                buf.expected_size = decode(&buf.buf[..]).unwrap();
                // Extend the Vec so we've got enough room for the remaning bytes, minus the 8 for the
                // encoded size header
                buf.buf.extend(iter::repeat(0u8).take(buf.expected_size - 8));
            }
        }
        // If we've finished reading the size header we can now start reading the frame data
        if buf.currently_read >= 8 {
            match self.connections[worker].read(&mut buf.buf[buf.currently_read..]) {
                Ok(n) => buf.currently_read += n,
                Err(e) => println!("Error reading results from worker {}: {}", self.workers[worker], e),
            }
        }
        buf.currently_read == buf.expected_size
    }
}

impl Handler for Master {
    type Timeout = ();
    type Message = ();

    fn ready(&mut self, event_loop: &mut EventLoop<Master>, token: Token, event: EventSet) {
        let worker = token.as_usize();
        if event.is_error() {
            println!("Error connecting too {}", self.workers[worker]);
            match self.connections[worker].shutdown(Shutdown::Both) {
                Err(e) => println!("Error shutting down worker {}: {}", worker, e),
                Ok(_) => {},
            }
            // Remove the connection from the event loop
            match event_loop.deregister(&self.connections[worker]) {
                Err(e) => println!("Error deregistering worker {}: {}", worker, e),
                Ok(_) => {},
            }
        }
        // If the worker has terminated, shutdown the read end of the connection
        if event.is_hup() {
            match self.connections[worker].shutdown(Shutdown::Both) {
                Err(e) => println!("Error shutting down worker {}: {}", worker, e),
                Ok(_) => {},
            }
            // Remove the connection from the event loop
            match event_loop.deregister(&self.connections[worker]) {
                Err(e) => println!("Error deregistering worker {}: {}", worker, e),
                Ok(_) => {},
            }
        }
        // A worker is ready to receive instructions from us
        if event.is_writable() {
            let b_start = worker * self.blocks_per_worker;
            let b_count =
                if worker == self.workers.len() - 1 {
                    self.blocks_per_worker + self.blocks_remainder
                } else {
                    self.blocks_per_worker
                };
            let instr = Instructions::new(&self.config.scene_file,
                                          (self.config.frame_info.start, self.config.frame_info.end),
                                          b_start, b_count);
            // Encode and send our instructions to the worker
            let bytes = encode(&instr, SizeLimit::Infinite).unwrap();
            match self.connections[worker].write_all(&bytes[..]) {
                Err(e) => println!("Failed to send instructions to {}: {:?}", self.workers[worker], e),
                Ok(_) => {},
            }
            // Register that we no longer care about writable events on this connection
            event_loop.reregister(&self.connections[worker], token,
                                  EventSet::readable() | EventSet::error() | EventSet::hup(),
                                  PollOpt::level()).expect("Re-registering failed");
        }
        // Some results are available from a worker
        if event.is_readable() {
            // Read results from the worker, if we've accumulated all the data being sent
            // decode and accumulate the frame
            if self.read_worker_buffer(worker) {
                let frame: Frame = decode(&self.worker_buffers[worker].buf[..]).unwrap();
                self.save_results(frame);
                // Clean up the worker buffer for the next frame
                self.worker_buffers[worker].buf.clear();
                self.worker_buffers[worker].expected_size = 8;
                self.worker_buffers[worker].currently_read = 0;
            }
        }
        // After getting results from the worker we check if we've completed all our frames
        // and exit if so
        let all_complete = self.frames.values().fold(true,
                                |all, v| {
                                    match v {
                                        &DistributedFrame::Completed => true && all,
                                        _ => false,
                                    }
                                });
        // The frame start/end range is inclusive, so we must add 1 here
        let num_frames = self.config.frame_info.end - self.config.frame_info.start + 1;
        if self.frames.len() == num_frames && all_complete {
            event_loop.shutdown();
        }
    }
}

<|endoftext|>


<commit_before>//! Given:
//!
//! `y <- alpha * A * x`
//!
//! Test that
//!
//! `y[i] = alpha * A[i, :] * x`
//!
//! for any valid `i`

#![feature(custom_attribute)]
#![feature(plugin)]
#![plugin(quickcheck_macros)]

extern crate approx;
extern crate complex;
extern crate linalg;
extern crate quickcheck;
extern crate rand;

#[macro_use]
mod setup;

use complex::{c64, c128};
use linalg::prelude::*;
use quickcheck::TestResult;

mod transposed {
    use complex::{c64, c128};
    use linalg::prelude::*;
    use quickcheck::TestResult;

    macro_rules! tests {
        ($($t:ident),+) => {
            $(
                #[quickcheck]
                fn $t((srow, scol): (u32, u32), (nrows, ncols): (u32, u32), i: u32) -> TestResult {
                    enforce! {
                        i < nrows,
                        ncols != 0,
                    }

                    let a = ::setup::rand::mat((srow + ncols, scol + nrows));
                    let a = a.slice((srow.., scol..)).t();
                    let ref x = ::setup::rand::col(ncols);
                    let alpha: $t = ::setup::rand::scalar();
                    let mut y = ::setup::rand::col(nrows);

                    y.set(alpha * a * x);

                    test_approx_eq!(y[i], alpha * a.row(i) * x)
                }
             )+
        }
    }

    tests!(f32, f64, c64, c128);

}

macro_rules! tests {
    ($($t:ident),+) => {
        $(
            #[quickcheck]
            fn $t((srow, scol): (u32, u32), (nrows, ncols): (u32, u32), i: u32) -> TestResult {
                enforce! {
                    i < nrows,
                    ncols != 0,
                }

                let a = ::setup::rand::mat((srow + nrows, scol + ncols));
                let a = a.slice((srow.., scol..));
                let ref x = ::setup::rand::col(ncols);
                let alpha: $t = ::setup::rand::scalar();
                let mut y = ::setup::rand::col(nrows);

                y.set(alpha * a * x);

                test_approx_eq!(y[i], alpha * a.row(i) * x)
            }
         )+
    }
}

tests!(f32, f64, c64, c128);

//mod col {
    //macro_rules! tests {
        //($($t:ident),+) => {
            //$(
                //#[quickcheck]
                //fn $t(
                    //(srow, scol): (u32, u32),
                    //(nrows, ncols): (u32, u32),
                    //(c1, c2, row): (u32, u32, u32),
                //) -> TestResult {
                    //enforce! {
                        //c1 < ncols,
                        //c2 < ncols,
                        //row < nrows,
                    //}

                    //let mut a = ::setup::rand::mat::<$t>((srow + nrows, scol + ncols));
                    //let mut a = a.slice_mut((srow.., scol..));
                    //let b = ::setup::rand::mat((srow + nrows, scol + ncols));
                    //let b = b.slice((srow.., scol..));

                    //a.col_mut(c1).set(b.col(c2));

                    //test_eq!(a[(row, c1)], b[(row, c2)])
                //}
             //)+
        //}
    //}

    //tests!(f32, f64, c64, c128);
//}

//mod row {
    //use complex::{c64, c128};
    //use linalg::prelude::*;
    //use quickcheck::TestResult;

    //macro_rules! tests {
        //($($t:ident),+) => {
            //$(
                //#[quickcheck]
                //fn $t(
                    //(srow, scol): (u32, u32),
                    //(nrows, ncols): (u32, u32),
                    //(r1, r2, col): (u32, u32, u32),
                //) -> TestResult {
                    //enforce! {
                        //r1 < nrows,
                        //r2 < nrows,
                        //col < ncols,
                    //}

                    //let mut a = ::setup::rand::mat::<$t>((srow + nrows, scol + ncols));
                    //let mut a = a.slice_mut((srow.., scol..));
                    //let b = ::setup::rand::mat((srow + nrows, scol + ncols));
                    //let b = b.slice((srow.., scol..));

                    //a.row_mut(r1).set(b.row(r2));

                    //test_eq!(a[(r1, col)], b[(r2, col)])
                //}
             //)+
        //}
    //}

    //tests!(f32, f64, c64, c128);
//}

//mod submat {
    //use complex::{c64, c128};
    //use linalg::prelude::*;
    //use quickcheck::TestResult;

    //macro_rules! tests {
        //($($t:ident),+) => {
            //$(
                //#[quickcheck]
                //fn $t(
                    //(srow, scol): (u32, u32),
                    //(nrows, ncols): (u32, u32),
                    //(row, col): (u32, u32),
                //) -> TestResult {
                    //enforce! {
                        //row < nrows,
                        //col < ncols,
                    //}

                    //let mut a = ::setup::rand::mat::<$t>((srow + nrows, scol + ncols));
                    //let mut a = a.slice_mut((srow.., scol..));
                    //let b = ::setup::rand::mat((srow + nrows, scol + ncols));
                    //let b = b.slice((srow.., scol..));
                    //a.set(b);

                    //test_eq!(a[(row, col)], b[(row, col)])
                //}
             //)+
        //}
    //}

    //tests!(f32, f64, c64, c128);
//}
<commit_msg>refactor: prune commented out code<commit_after>//! Given:
//!
//! `y <- alpha * A * x`
//!
//! Test that
//!
//! `y[i] = alpha * A[i, :] * x`
//!
//! for any valid `i`

#![feature(custom_attribute)]
#![feature(plugin)]
#![plugin(quickcheck_macros)]

extern crate approx;
extern crate complex;
extern crate linalg;
extern crate quickcheck;
extern crate rand;

#[macro_use]
mod setup;

use complex::{c64, c128};
use linalg::prelude::*;
use quickcheck::TestResult;

mod transposed {
    use complex::{c64, c128};
    use linalg::prelude::*;
    use quickcheck::TestResult;

    macro_rules! tests {
        ($($t:ident),+) => {
            $(
                #[quickcheck]
                fn $t((srow, scol): (u32, u32), (nrows, ncols): (u32, u32), i: u32) -> TestResult {
                    enforce! {
                        i < nrows,
                        ncols != 0,
                    }

                    let a = ::setup::rand::mat((srow + ncols, scol + nrows));
                    let a = a.slice((srow.., scol..)).t();
                    let ref x = ::setup::rand::col(ncols);
                    let alpha: $t = ::setup::rand::scalar();
                    let mut y = ::setup::rand::col(nrows);

                    y.set(alpha * a * x);

                    test_approx_eq!(y[i], alpha * a.row(i) * x)
                }
             )+
        }
    }

    tests!(f32, f64, c64, c128);

}

macro_rules! tests {
    ($($t:ident),+) => {
        $(
            #[quickcheck]
            fn $t((srow, scol): (u32, u32), (nrows, ncols): (u32, u32), i: u32) -> TestResult {
                enforce! {
                    i < nrows,
                    ncols != 0,
                }

                let a = ::setup::rand::mat((srow + nrows, scol + ncols));
                let a = a.slice((srow.., scol..));
                let ref x = ::setup::rand::col(ncols);
                let alpha: $t = ::setup::rand::scalar();
                let mut y = ::setup::rand::col(nrows);

                y.set(alpha * a * x);

                test_approx_eq!(y[i], alpha * a.row(i) * x)
            }
         )+
    }
}

tests!(f32, f64, c64, c128);
<|endoftext|>


<commit_before>use average::Histogram;
use rand::{Rng, SeedableRng};
use rand_distr::Normal;

average::define_histogram!(hist, 100);
use hist::Histogram as Histogram100;

mod sparkline;

#[test]
fn normal() {
    const N_SAMPLES: u64 = 1_000_000;
    const MEAN: f64 = 2.;
    const STD_DEV: f64 = 0.5;
    const MIN_X: f64 = -1.;
    const MAX_X: f64 = 5.;

    let dist = Normal::new(MEAN, STD_DEV).unwrap();
    let mut hist = Histogram100::with_const_width(MIN_X,MAX_X);
    let mut rng = rand::rngs::SmallRng::seed_from_u64(1);

    for _ in 0..N_SAMPLES {
        let _ = hist.add(rng.sample(dist));  // Ignore out-of-range values
    }

    println!("Sampled normal distribution:\n{}",
        sparkline::render_u64_as_string(hist.bins()));

    fn pdf(x: f64) -> f64 {
        (-0.5 * ((x - MEAN) / STD_DEV).powi(2)).exp() /
            (STD_DEV * (2. * core::f64::consts::PI).sqrt())
    }

    let mut bin_centers = hist.centers();
    let mut expected = [0.; 100];
    for e in &mut expected {
        *e = pdf(bin_centers.next().unwrap());
    }

    println!("Expected normal distribution:\n{}",
        sparkline::render_u64_as_string(hist.bins()));

    let mut normalized_bins= hist.normalized_bins();
    let mut diff = [0.; 100];
    for i in 0..100 {
        let bin = (normalized_bins.next().unwrap() as f64) / (N_SAMPLES as f64) ;
        diff[i] = (bin - expected[i]).abs();
    }

    println!("Difference:\n{}",
        sparkline::render_f64_as_string(&diff[..]));
    println!("max diff: {:?}", diff.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b)));

    // Check that the differences are significantly smaller than the expected error.
    let mut expected_error = [0.; 100];
    // Calculate error from histogram
    for (err, var) in expected_error.iter_mut().zip(hist.variances()) {
        *err = var.sqrt() / (N_SAMPLES as f64);
    }
    // Normalize error by bin width
    for (err, width) in expected_error.iter_mut().zip(hist.widths()) {
        *err /= width;
    }
    // TODO: Calculate error from distribution cutoff / normalization

    println!("max expected_error: {:?}", expected_error.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b)));
    for (&d, &e) in diff.iter().zip(expected_error.iter()) {
        // Difference larger than 3 standard deviations or cutoff
        let tol = (3. * e).max(1e-4);
        if d > tol {
            panic!("Difference = {} * tol", d / tol);
        }
    }
}<commit_msg>Fix tests for old Rust versions<commit_after>use average::Histogram;
use rand::{Rng, SeedableRng};
use rand_distr::Normal;

average::define_histogram!(hist, 100);
use hist::Histogram as Histogram100;

mod sparkline;

#[test]
fn normal() {
    const N_SAMPLES: u64 = 1_000_000;
    const MEAN: f64 = 2.;
    const STD_DEV: f64 = 0.5;
    const MIN_X: f64 = -1.;
    const MAX_X: f64 = 5.;

    let dist = Normal::new(MEAN, STD_DEV).unwrap();
    let mut hist = Histogram100::with_const_width(MIN_X,MAX_X);
    let mut rng = rand::rngs::SmallRng::seed_from_u64(1);

    for _ in 0..N_SAMPLES {
        let _ = hist.add(rng.sample(dist));  // Ignore out-of-range values
    }

    println!("Sampled normal distribution:\n{}",
        sparkline::render_u64_as_string(hist.bins()));

    fn pdf(x: f64) -> f64 {
        (-0.5 * ((x - MEAN) / STD_DEV).powi(2)).exp() /
            (STD_DEV * (2. * core::f64::consts::PI).sqrt())
    }

    let mut bin_centers = hist.centers();
    let mut expected = [0.; 100];
    for e in &mut expected {
        *e = pdf(bin_centers.next().unwrap());
    }

    println!("Expected normal distribution:\n{}",
        sparkline::render_u64_as_string(hist.bins()));

    let mut normalized_bins= hist.normalized_bins();
    let mut diff = [0.; 100];
    for i in 0..100 {
        let bin = (normalized_bins.next().unwrap() as f64) / (N_SAMPLES as f64) ;
        diff[i] = (bin - expected[i]).abs();
    }

    println!("Difference:\n{}",
        sparkline::render_f64_as_string(&diff[..]));
    println!("max diff: {:?}", diff.iter().fold(
        core::f64::NEG_INFINITY, |a, &b| a.max(b)));

    // Check that the differences are significantly smaller than the expected error.
    let mut expected_error = [0.; 100];
    // Calculate error from histogram
    for (err, var) in expected_error.iter_mut().zip(hist.variances()) {
        *err = var.sqrt() / (N_SAMPLES as f64);
    }
    // Normalize error by bin width
    for (err, width) in expected_error.iter_mut().zip(hist.widths()) {
        *err /= width;
    }
    // TODO: Calculate error from distribution cutoff / normalization

    println!("max expected_error: {:?}", expected_error.iter().fold(
        core::f64::NEG_INFINITY, |a, &b| a.max(b)));
    for (&d, &e) in diff.iter().zip(expected_error.iter()) {
        // Difference larger than 3 standard deviations or cutoff
        let tol = (3. * e).max(1e-4);
        if d > tol {
            panic!("Difference = {} * tol", d / tol);
        }
    }
}<|endoftext|>