<commit_before>use super::ipv4_header::{IPv4Header, IPv4HeaderData, IPv4HeaderMut};
use super::transport_header::{TransportHeader, TransportHeaderData, TransportHeaderMut};

pub const MAX_PACKET_LENGTH: usize = 1 << 16;

pub struct IPv4Packet<'a> {
    raw: &'a mut [u8],
    ipv4_header_data: IPv4HeaderData,
    transport_header_data: Option<TransportHeaderData>,
}

impl<'a> IPv4Packet<'a> {
    pub fn parse(raw: &'a mut [u8]) -> Self {
        let ipv4_header_data = IPv4HeaderData::parse(raw);
        let transport_header_data = {
            let payload = &raw[ipv4_header_data.header_length() as usize..];
            TransportHeaderData::parse(ipv4_header_data.protocol(), payload)
        };
        Self {
            raw: &mut raw[..ipv4_header_data.total_length() as usize],
            ipv4_header_data: ipv4_header_data,
            transport_header_data: transport_header_data,
        }
    }

    pub fn new(raw: &'a mut [u8], ipv4_header_data: IPv4HeaderData, transport_header_data: TransportHeaderData) -> Self {
        Self {
            raw: raw,
            ipv4_header_data: ipv4_header_data,
            transport_header_data: Some(transport_header_data),
        }
    }

    #[inline]
    pub fn raw(&self) -> &[u8] {
        self.raw
    }

    #[inline]
    pub fn raw_mut(&mut self) -> &mut [u8] {
        self.raw
    }

    #[inline]
    pub fn ipv4_header_data(&self) -> &IPv4HeaderData {
        &self.ipv4_header_data
    }

    #[inline]
    pub fn ipv4_header(&self) -> IPv4Header {
        let slice = &self.raw[..self.ipv4_header_data.header_length() as usize];
        self.ipv4_header_data.bind(slice)
    }

    #[inline]
    pub fn ipv4_header_mut(&mut self) -> IPv4HeaderMut {
        let slice = &mut self.raw[..self.ipv4_header_data.header_length() as usize];
        self.ipv4_header_data.bind_mut(slice)
    }

    #[inline]
    pub fn transport_header_data(&self) -> &Option<TransportHeaderData> {
        &self.transport_header_data
    }

    #[inline]
    pub fn transport_header(&self) -> Option<TransportHeader> {
        if let Some(ref transport_header_data) = self.transport_header_data {
            let start = self.ipv4_header_data.header_length() as usize;
            let end = start + transport_header_data.header_length() as usize;
            let slice = &self.raw[start..end];
            Some(transport_header_data.bind(slice))
        } else {
            None
        }
/*        self.transport_header_data.as_ref().map(|transport_header_data| {
            let start = self.ipv4_header_data.header_length() as usize;
            let end = start + transport_header_data.header_length() as usize;
            let slice = &self.raw[start..end];
            transport_header_data.bind(slice)
        })*/
    }

    #[inline]
    pub fn transport_header_mut(&mut self) -> Option<TransportHeaderMut> {
        if let Some(ref mut transport_header_data) = self.transport_header_data {
            let start = self.ipv4_header_data.header_length() as usize;
            let end = start + transport_header_data.header_length() as usize;
            let slice = &mut self.raw[start..end];
            Some(transport_header_data.bind_mut(slice))
        } else {
            None
        }
/*        self.transport_header_data.as_mut().map(|transport_header_data| {
            let start = self.ipv4_header_data.header_length() as usize;
            let end = start + transport_header_data.header_length() as usize;
            let slice = &mut self.raw[start..end];
            transport_header_data.bind_mut(slice)
        })*/
    }

    /// Devide the packet into parts:
    ///  - the IPv4 header
    ///  - the transport header (if any)
    ///  - the payload (if there is a transport at all)
    pub fn split(&self) -> (IPv4Header, Option<(TransportHeader, &[u8])>) {
        let transport_index = self.ipv4_header_data.header_length() as usize;
        if let Some(ref transport_header_data) = self.transport_header_data {
            let payload_index = transport_header_data.header_length() as usize; // relative to transport
            let (ipv4_header_slice, transport_slice) = self.raw.split_at(transport_index);
            let (transport_header_slice, payload_slice) = transport_slice.split_at(payload_index);
            let ipv4_header = self.ipv4_header_data.bind(ipv4_header_slice);
            let transport_header = transport_header_data.bind(transport_header_slice);
            (ipv4_header, Some((transport_header, payload_slice)))
        } else {
            let ipv4_header_slice = &self.raw[..transport_index];
            let ipv4_header = self.ipv4_header_data.bind(ipv4_header_slice);
            (ipv4_header, None)
        }
    }

    /// Devide the packet into mutable parts:
    ///  - the IPv4 header
    ///  - the transport header (if any)
    ///  - the payload (if there is a transport at all)
    pub fn split_mut(&mut self) -> (IPv4HeaderMut, Option<(TransportHeaderMut, &mut [u8])>) {
        let transport_index = self.ipv4_header_data.header_length() as usize;
        if let Some(ref mut transport_header_data) = self.transport_header_data {
            let payload_index = transport_header_data.header_length() as usize; // relative to transport
            let (ipv4_header_slice, transport_slice) = self.raw.split_at_mut(transport_index);
            let (transport_header_slice, payload_slice) = transport_slice.split_at_mut(payload_index);
            let ipv4_header = self.ipv4_header_data.bind_mut(ipv4_header_slice);
            let transport_header = transport_header_data.bind_mut(transport_header_slice);
            (ipv4_header, Some((transport_header, payload_slice)))
        } else {
            let ipv4_header_slice = &mut self.raw[..transport_index];
            let ipv4_header = self.ipv4_header_data.bind_mut(ipv4_header_slice);
            (ipv4_header, None)
        }
    }

    #[inline]
    pub fn is_valid(&self) -> bool {
        self.transport_header_data.is_some()
    }

    #[inline]
    pub fn length(&self) -> u16 {
        self.ipv4_header_data.total_length()
    }

/*
    pub fn ipv4_header_range(&self) -> Range<usize> {
        let start = 0;
        let end = self.ipv4_header.header_length() as usize;
        start..end
    }

    pub fn transport_range(&self) -> Option<Range<usize>> {
        self.transport_header.as_ref().map(|_| {
            let start = self.ipv4_header.header_length() as usize;
            let end = self.raw.len();
            start..end
        })
    }

    pub fn transport_header_range(&self) -> Option<Range<usize>> {
        self.transport_header.as_ref().map(|transport_header| {
            let start = self.ipv4_header.header_length() as usize;
            let end = start + transport_header.header_length() as usize;
            start..end
        })
    }

    pub fn payload_range(&self) -> Option<Range<usize>> {
        self.transport_header.as_ref().map(|transport_header| {
            let start = self.ipv4_header.header_length() as usize + transport_header.header_length() as usize;
            let end = self.raw.len();
            start..end
        })
    }
*/
    // TODO delete function
    pub fn payload(&self) -> Option<&[u8]> {
        self.transport_header_data.as_ref().map(|transport_header_data| {
            let range = self.ipv4_header_data.header_length() as usize + transport_header_data.header_length() as usize..;
            &self.raw[range]
        })
    }

    pub fn compute_checksums(&mut self) {
        let (mut ipv4_header, transport) = self.split_mut();
        ipv4_header.update_checksum();
        if let Some((mut transport_header, payload)) = transport {
            transport_header.update_checksum(ipv4_header.data(), payload);
        }
    }

    #[inline]
    pub fn swap_source_and_destination(&mut self) {
        self.ipv4_header_mut().swap_source_and_destination();
        if let Some(mut transport_header) = self.transport_header_mut() {
            transport_header.swap_source_and_destination();
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use byteorder::{BigEndian, ByteOrder, WriteBytesExt};
    use ::relay::ipv4_header::Protocol;

    fn create_packet() -> Vec<u8> {
        let mut raw = Vec::new();
        raw.reserve(32);

        raw.write_u8(4u8 << 4 | 5).unwrap(); // version_and_ihl
        raw.write_u8(0).unwrap(); //ToS
        raw.write_u16::<BigEndian>(32).unwrap(); // total length 20 + 8 + 4
        raw.write_u32::<BigEndian>(0).unwrap(); // id_flags_fragment_offset
        raw.write_u8(0).unwrap(); // TTL
        raw.write_u8(17).unwrap(); // protocol (UDP)
        raw.write_u16::<BigEndian>(0).unwrap(); // checksum
        raw.write_u32::<BigEndian>(0x12345678).unwrap(); // source address
        raw.write_u32::<BigEndian>(0x42424242).unwrap(); // destination address

        raw.write_u16::<BigEndian>(1234).unwrap(); // source port
        raw.write_u16::<BigEndian>(5678).unwrap(); // destination port
        raw.write_u16::<BigEndian>(4).unwrap(); // length
        raw.write_u16::<BigEndian>(0).unwrap(); // checksum

        raw.write_u32::<BigEndian>(0x11223344).unwrap(); // payload

        raw
    }

    #[test]
    fn parse_headers() {
        let raw = &mut create_packet()[..];
        let mut ipv4_packet = IPv4Packet::parse(raw);

        {
            let ipv4_header = ipv4_packet.ipv4_header();
            assert_eq!(20, ipv4_header.header_length());
            assert_eq!(32, ipv4_header.total_length());
            assert_eq!(Protocol::UDP, ipv4_header.protocol());
            assert_eq!(0x12345678, ipv4_header.source());
            assert_eq!(0x42424242, ipv4_header.destination());

            if let Some(TransportHeaderData::UDP(ref udp_header)) = *ipv4_packet.transport_header_data() {
                assert_eq!(1234, udp_header.source_port());
                assert_eq!(5678, udp_header.destination_port());
            } else {
                panic!("No UDP transport header");
            }
        }

        ipv4_packet.swap_source_and_destination();

        {
            let ipv4_header = ipv4_packet.ipv4_header();
            assert_eq!(0x42424242, ipv4_header.source());
            assert_eq!(0x12345678, ipv4_header.destination());

            if let Some(TransportHeaderData::UDP(ref udp_header)) = *ipv4_packet.transport_header_data() {
                assert_eq!(5678, udp_header.source_port());
                assert_eq!(1234, udp_header.destination_port());
            } else {
                panic!("No UDP transport header");
            }
        }

        {
            let raw = ipv4_packet.raw();
            // assert that the buffer has been modified
            let raw_source = BigEndian::read_u32(&raw[12..16]);
            let raw_destination = BigEndian::read_u32(&raw[16..20]);
            assert_eq!(0x42424242, raw_source);
            assert_eq!(0x12345678, raw_destination);

            if let Some(TransportHeaderData::UDP(ref udp_header)) = *ipv4_packet.transport_header_data() {
                assert_eq!(5678, udp_header.source_port());
                assert_eq!(1234, udp_header.destination_port());
            } else {
                panic!("No UDP transport header");
            }
        }
    }

    #[test]
    fn payload() {
        let raw = &mut create_packet()[..];
        let ipv4_packet = IPv4Packet::parse(raw);
        assert_eq!([0x11, 0x22, 0x33, 0x44], ipv4_packet.payload().unwrap());
    }
}
<commit_msg>Delete old commented code<commit_after>use super::ipv4_header::{IPv4Header, IPv4HeaderData, IPv4HeaderMut};
use super::transport_header::{TransportHeader, TransportHeaderData, TransportHeaderMut};

pub const MAX_PACKET_LENGTH: usize = 1 << 16;

pub struct IPv4Packet<'a> {
    raw: &'a mut [u8],
    ipv4_header_data: IPv4HeaderData,
    transport_header_data: Option<TransportHeaderData>,
}

impl<'a> IPv4Packet<'a> {
    pub fn parse(raw: &'a mut [u8]) -> Self {
        let ipv4_header_data = IPv4HeaderData::parse(raw);
        let transport_header_data = {
            let payload = &raw[ipv4_header_data.header_length() as usize..];
            TransportHeaderData::parse(ipv4_header_data.protocol(), payload)
        };
        Self {
            raw: &mut raw[..ipv4_header_data.total_length() as usize],
            ipv4_header_data: ipv4_header_data,
            transport_header_data: transport_header_data,
        }
    }

    pub fn new(raw: &'a mut [u8], ipv4_header_data: IPv4HeaderData, transport_header_data: TransportHeaderData) -> Self {
        Self {
            raw: raw,
            ipv4_header_data: ipv4_header_data,
            transport_header_data: Some(transport_header_data),
        }
    }

    #[inline]
    pub fn raw(&self) -> &[u8] {
        self.raw
    }

    #[inline]
    pub fn raw_mut(&mut self) -> &mut [u8] {
        self.raw
    }

    #[inline]
    pub fn ipv4_header_data(&self) -> &IPv4HeaderData {
        &self.ipv4_header_data
    }

    #[inline]
    pub fn ipv4_header(&self) -> IPv4Header {
        let slice = &self.raw[..self.ipv4_header_data.header_length() as usize];
        self.ipv4_header_data.bind(slice)
    }

    #[inline]
    pub fn ipv4_header_mut(&mut self) -> IPv4HeaderMut {
        let slice = &mut self.raw[..self.ipv4_header_data.header_length() as usize];
        self.ipv4_header_data.bind_mut(slice)
    }

    #[inline]
    pub fn transport_header_data(&self) -> &Option<TransportHeaderData> {
        &self.transport_header_data
    }

    #[inline]
    pub fn transport_header(&self) -> Option<TransportHeader> {
        if let Some(ref transport_header_data) = self.transport_header_data {
            let start = self.ipv4_header_data.header_length() as usize;
            let end = start + transport_header_data.header_length() as usize;
            let slice = &self.raw[start..end];
            Some(transport_header_data.bind(slice))
        } else {
            None
        }
/*        self.transport_header_data.as_ref().map(|transport_header_data| {
            let start = self.ipv4_header_data.header_length() as usize;
            let end = start + transport_header_data.header_length() as usize;
            let slice = &self.raw[start..end];
            transport_header_data.bind(slice)
        })*/
    }

    #[inline]
    pub fn transport_header_mut(&mut self) -> Option<TransportHeaderMut> {
        if let Some(ref mut transport_header_data) = self.transport_header_data {
            let start = self.ipv4_header_data.header_length() as usize;
            let end = start + transport_header_data.header_length() as usize;
            let slice = &mut self.raw[start..end];
            Some(transport_header_data.bind_mut(slice))
        } else {
            None
        }
/*        self.transport_header_data.as_mut().map(|transport_header_data| {
            let start = self.ipv4_header_data.header_length() as usize;
            let end = start + transport_header_data.header_length() as usize;
            let slice = &mut self.raw[start..end];
            transport_header_data.bind_mut(slice)
        })*/
    }

    /// Devide the packet into parts:
    ///  - the IPv4 header
    ///  - the transport header (if any)
    ///  - the payload (if there is a transport at all)
    pub fn split(&self) -> (IPv4Header, Option<(TransportHeader, &[u8])>) {
        let transport_index = self.ipv4_header_data.header_length() as usize;
        if let Some(ref transport_header_data) = self.transport_header_data {
            let payload_index = transport_header_data.header_length() as usize; // relative to transport
            let (ipv4_header_slice, transport_slice) = self.raw.split_at(transport_index);
            let (transport_header_slice, payload_slice) = transport_slice.split_at(payload_index);
            let ipv4_header = self.ipv4_header_data.bind(ipv4_header_slice);
            let transport_header = transport_header_data.bind(transport_header_slice);
            (ipv4_header, Some((transport_header, payload_slice)))
        } else {
            let ipv4_header_slice = &self.raw[..transport_index];
            let ipv4_header = self.ipv4_header_data.bind(ipv4_header_slice);
            (ipv4_header, None)
        }
    }

    /// Devide the packet into mutable parts:
    ///  - the IPv4 header
    ///  - the transport header (if any)
    ///  - the payload (if there is a transport at all)
    pub fn split_mut(&mut self) -> (IPv4HeaderMut, Option<(TransportHeaderMut, &mut [u8])>) {
        let transport_index = self.ipv4_header_data.header_length() as usize;
        if let Some(ref mut transport_header_data) = self.transport_header_data {
            let payload_index = transport_header_data.header_length() as usize; // relative to transport
            let (ipv4_header_slice, transport_slice) = self.raw.split_at_mut(transport_index);
            let (transport_header_slice, payload_slice) = transport_slice.split_at_mut(payload_index);
            let ipv4_header = self.ipv4_header_data.bind_mut(ipv4_header_slice);
            let transport_header = transport_header_data.bind_mut(transport_header_slice);
            (ipv4_header, Some((transport_header, payload_slice)))
        } else {
            let ipv4_header_slice = &mut self.raw[..transport_index];
            let ipv4_header = self.ipv4_header_data.bind_mut(ipv4_header_slice);
            (ipv4_header, None)
        }
    }

    #[inline]
    pub fn is_valid(&self) -> bool {
        self.transport_header_data.is_some()
    }

    #[inline]
    pub fn length(&self) -> u16 {
        self.ipv4_header_data.total_length()
    }

    pub fn payload(&self) -> Option<&[u8]> {
        self.transport_header_data.as_ref().map(|transport_header_data| {
            let range = self.ipv4_header_data.header_length() as usize + transport_header_data.header_length() as usize..;
            &self.raw[range]
        })
    }

    pub fn compute_checksums(&mut self) {
        let (mut ipv4_header, transport) = self.split_mut();
        ipv4_header.update_checksum();
        if let Some((mut transport_header, payload)) = transport {
            transport_header.update_checksum(ipv4_header.data(), payload);
        }
    }

    #[inline]
    pub fn swap_source_and_destination(&mut self) {
        self.ipv4_header_mut().swap_source_and_destination();
        if let Some(mut transport_header) = self.transport_header_mut() {
            transport_header.swap_source_and_destination();
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use byteorder::{BigEndian, ByteOrder, WriteBytesExt};
    use ::relay::ipv4_header::Protocol;

    fn create_packet() -> Vec<u8> {
        let mut raw = Vec::new();
        raw.reserve(32);

        raw.write_u8(4u8 << 4 | 5).unwrap(); // version_and_ihl
        raw.write_u8(0).unwrap(); //ToS
        raw.write_u16::<BigEndian>(32).unwrap(); // total length 20 + 8 + 4
        raw.write_u32::<BigEndian>(0).unwrap(); // id_flags_fragment_offset
        raw.write_u8(0).unwrap(); // TTL
        raw.write_u8(17).unwrap(); // protocol (UDP)
        raw.write_u16::<BigEndian>(0).unwrap(); // checksum
        raw.write_u32::<BigEndian>(0x12345678).unwrap(); // source address
        raw.write_u32::<BigEndian>(0x42424242).unwrap(); // destination address

        raw.write_u16::<BigEndian>(1234).unwrap(); // source port
        raw.write_u16::<BigEndian>(5678).unwrap(); // destination port
        raw.write_u16::<BigEndian>(4).unwrap(); // length
        raw.write_u16::<BigEndian>(0).unwrap(); // checksum

        raw.write_u32::<BigEndian>(0x11223344).unwrap(); // payload

        raw
    }

    #[test]
    fn parse_headers() {
        let raw = &mut create_packet()[..];
        let mut ipv4_packet = IPv4Packet::parse(raw);

        {
            let ipv4_header = ipv4_packet.ipv4_header();
            assert_eq!(20, ipv4_header.header_length());
            assert_eq!(32, ipv4_header.total_length());
            assert_eq!(Protocol::UDP, ipv4_header.protocol());
            assert_eq!(0x12345678, ipv4_header.source());
            assert_eq!(0x42424242, ipv4_header.destination());

            if let Some(TransportHeaderData::UDP(ref udp_header)) = *ipv4_packet.transport_header_data() {
                assert_eq!(1234, udp_header.source_port());
                assert_eq!(5678, udp_header.destination_port());
            } else {
                panic!("No UDP transport header");
            }
        }

        ipv4_packet.swap_source_and_destination();

        {
            let ipv4_header = ipv4_packet.ipv4_header();
            assert_eq!(0x42424242, ipv4_header.source());
            assert_eq!(0x12345678, ipv4_header.destination());

            if let Some(TransportHeaderData::UDP(ref udp_header)) = *ipv4_packet.transport_header_data() {
                assert_eq!(5678, udp_header.source_port());
                assert_eq!(1234, udp_header.destination_port());
            } else {
                panic!("No UDP transport header");
            }
        }

        {
            let raw = ipv4_packet.raw();
            // assert that the buffer has been modified
            let raw_source = BigEndian::read_u32(&raw[12..16]);
            let raw_destination = BigEndian::read_u32(&raw[16..20]);
            assert_eq!(0x42424242, raw_source);
            assert_eq!(0x12345678, raw_destination);

            if let Some(TransportHeaderData::UDP(ref udp_header)) = *ipv4_packet.transport_header_data() {
                assert_eq!(5678, udp_header.source_port());
                assert_eq!(1234, udp_header.destination_port());
            } else {
                panic!("No UDP transport header");
            }
        }
    }

    #[test]
    fn payload() {
        let raw = &mut create_packet()[..];
        let ipv4_packet = IPv4Packet::parse(raw);
        assert_eq!([0x11, 0x22, 0x33, 0x44], ipv4_packet.payload().unwrap());
    }
}
<|endoftext|>


<commit_before>use linalg::matrix::Matrix;
use linalg::vector::Vector;
use learning::SupModel;
use rand::{Rng, thread_rng};

/// Neural Network struct
///
/// Requires the number of neurons in each layer to be specified.
pub struct NeuralNet<'a> {
    layer_sizes: &'a [usize],
    weights: Vec<f64>,
}

impl<'a> NeuralNet<'a> {

	/// Create a new neural network with the specified layer sizes.
	///
	/// The layer sizes slice should include the input, hidden layers, and output layer sizes.
	///
	/// # Examples
	///
	/// ```
	/// use rusty_machine::learning::nnet::NeuralNet;
	///
	/// // Create a neural net with 4 layers, 3 neurons in each.
	/// let layers = &[3; 4];
	/// let mut a = NeuralNet::new(layers);
	/// ```
    pub fn new(layer_sizes: &[usize]) -> NeuralNet {
        NeuralNet {
            layer_sizes: layer_sizes,
            weights: NeuralNet::create_weights(layer_sizes),
        }
    }

    /// Creates initial weights for all neurons in the network.
    fn create_weights(layer_sizes: &[usize]) -> Vec<f64> {
        let total_layers = layer_sizes.len();

        let mut capacity = 0usize;

        for l in 0..total_layers - 1 {
            capacity += (layer_sizes[l]+1) * layer_sizes[l + 1]
        }

        let mut layers = Vec::with_capacity(capacity);

        for l in 0..total_layers - 1 {
            layers.append(&mut NeuralNet::initialize_weights(layer_sizes[l]+1, layer_sizes[l + 1]));
        }

        layers
    }

    /// Initializes the weights for a single layer in the network.
    fn initialize_weights(rows: usize, cols: usize) -> Vec<f64> {
        let mut weights = Vec::with_capacity(rows * cols);
        let eps_init = (6f64).sqrt() / ((rows + cols) as f64).sqrt();

        let mut rng = thread_rng();

        for _i in 0..rows * cols {
            let w = (rng.gen_range(0f64, 1f64) * 2f64 * eps_init) - eps_init;
            weights.push(w);
        }

        weights
    }

    /// Gets matrix of weights between specified layer and forward layer.
	///
	/// # Examples
	///
	/// ```
	/// use rusty_machine::learning::nnet::NeuralNet;
	///
	/// // Create a neural net with 4 layers, 3 neurons in each.
	/// let layers = &[3; 4];
	/// let mut a = NeuralNet::new(layers);
	///
	/// let w = &a.get_layer_weights(2);
	/// assert_eq!(w.rows(), 4);
	/// assert_eq!(w.cols(), 3);
	/// ```
    pub fn get_layer_weights(&self, idx: usize) -> Matrix<f64> {
    	assert!(idx < self.layer_sizes.len()-1);

    	let mut start = 0usize;

        for l in 0..idx {
            start += (self.layer_sizes[l]+1) * self.layer_sizes[l + 1]
        }

        let capacity = (self.layer_sizes[idx]+1) * self.layer_sizes[idx + 1];

        let mut weights = Vec::with_capacity((self.layer_sizes[idx]+1) * self.layer_sizes[idx + 1]);
        unsafe {
	        for i in start..start+capacity {
	        	weights.push(*self.weights.get_unchecked(i));
	        }
    	}

        Matrix::new(self.layer_sizes[idx]+1, self.layer_sizes[idx + 1], weights)

    }

    fn compute_cost_grad(&self, data: &Matrix<f64>, outputs: &Matrix<f64>) {
    	assert_eq!(data.cols(), self.layer_sizes[0]);

    	let mut forward_weights = Vec::with_capacity(self.layer_sizes.len()-1);
    	let mut activations = Vec::with_capacity(self.layer_sizes.len());

    	activations.push(data.clone());

    	// Forward propagation
    	{
	    	let mut a = self.get_layer_weights(0) * data;
	    	activations.push(a.clone());

	    	for l in 1..self.layer_sizes.len()-1 {
	    		let z = self.get_layer_weights(l) * a.clone();
	    		forward_weights.push(z.clone());

	    		activations.push(z.apply(&sigmoid));
	    	}

	    }

	    let mut deltas = Vec::with_capacity(self.layer_sizes.len()-1);
	    // Backward propagation
	    {
	    	let mut delta = &activations[self.layer_sizes.len()-1] - outputs;
	    	deltas.push(delta.clone());

	    	for l in (0..self.layer_sizes.len()-1).rev() {
	    		let g = forward_weights[l].clone().apply(&sigmoid_grad);
	    		delta = (self.get_layer_weights(l).transpose() * delta.clone()).elemul(&g);
	    		deltas.push(delta.clone());
	    	}
	    }

	    let mut grad = Vec::with_capacity(self.layer_sizes.len()-1);
	   	for l in 0..self.layer_sizes.len()-1 {
	   		// Probably need to remove the first column on the deltas.
	   		let g = deltas[l].clone() * activations[l].clone().transpose();
	   		grad.push(g / (data.rows() as f64))
	   	}
    }

}

impl<'a> SupModel<Matrix<f64>, Vector<usize>> for NeuralNet<'a> {
    fn predict(&self, data: Matrix<f64>) -> Vector<usize> {
        unimplemented!()
    }

    fn train(&mut self, data: Matrix<f64>, values: Vector<usize>) {
        unimplemented!()
    }
}

/// Sigmoid function.
///
/// Returns 1 / ( 1 + e^-t).
fn sigmoid(t: f64) -> f64 {
	1.0 / (1.0 + (-t).exp() )
}

/// Gradient of sigmoid function.
///
/// Evaluates to (1 - e^-t) / (1 + e^-t)^2
fn sigmoid_grad(t: f64) -> f64 {
	sigmoid(t) * (1f64-sigmoid(t))
}<commit_msg>Implemented gradient computation for neural nets<commit_after>use linalg::matrix::Matrix;
use linalg::vector::Vector;
use learning::SupModel;
use rand::{Rng, thread_rng};

/// Neural Network struct
///
/// Requires the number of neurons in each layer to be specified.
pub struct NeuralNet<'a> {
    layer_sizes: &'a [usize],
    weights: Vec<f64>,
}

impl<'a> NeuralNet<'a> {

	/// Create a new neural network with the specified layer sizes.
	///
	/// The layer sizes slice should include the input, hidden layers, and output layer sizes.
	///
	/// # Examples
	///
	/// ```
	/// use rusty_machine::learning::nnet::NeuralNet;
	///
	/// // Create a neural net with 4 layers, 3 neurons in each.
	/// let layers = &[3; 4];
	/// let mut a = NeuralNet::new(layers);
	/// ```
    pub fn new(layer_sizes: &[usize]) -> NeuralNet {
        NeuralNet {
            layer_sizes: layer_sizes,
            weights: NeuralNet::create_weights(layer_sizes),
        }
    }

    /// Creates initial weights for all neurons in the network.
    fn create_weights(layer_sizes: &[usize]) -> Vec<f64> {
        let total_layers = layer_sizes.len();

        let mut capacity = 0usize;

        for l in 0..total_layers - 1 {
            capacity += (layer_sizes[l]+1) * layer_sizes[l + 1]
        }

        let mut layers = Vec::with_capacity(capacity);

        for l in 0..total_layers - 1 {
            layers.append(&mut NeuralNet::initialize_weights(layer_sizes[l]+1, layer_sizes[l + 1]));
        }

        layers
    }

    /// Initializes the weights for a single layer in the network.
    fn initialize_weights(rows: usize, cols: usize) -> Vec<f64> {
        let mut weights = Vec::with_capacity(rows * cols);
        let eps_init = (6f64).sqrt() / ((rows + cols) as f64).sqrt();

        let mut rng = thread_rng();

        for _i in 0..rows * cols {
            let w = (rng.gen_range(0f64, 1f64) * 2f64 * eps_init) - eps_init;
            weights.push(w);
        }

        weights
    }

    /// Gets matrix of weights between specified layer and forward layer.
	///
	/// # Examples
	///
	/// ```
	/// use rusty_machine::learning::nnet::NeuralNet;
	///
	/// // Create a neural net with 4 layers, 3 neurons in each.
	/// let layers = &[3; 4];
	/// let mut a = NeuralNet::new(layers);
	///
	/// let w = &a.get_layer_weights(2);
	/// assert_eq!(w.rows(), 4);
	/// assert_eq!(w.cols(), 3);
	/// ```
    pub fn get_layer_weights(&self, idx: usize) -> Matrix<f64> {
    	assert!(idx < self.layer_sizes.len()-1);

    	let mut start = 0usize;

        for l in 0..idx {
            start += (self.layer_sizes[l]+1) * self.layer_sizes[l + 1]
        }

        let capacity = (self.layer_sizes[idx]+1) * self.layer_sizes[idx + 1];

        let mut weights = Vec::with_capacity((self.layer_sizes[idx]+1) * self.layer_sizes[idx + 1]);
        unsafe {
	        for i in start..start+capacity {
	        	weights.push(*self.weights.get_unchecked(i));
	        }
    	}

        Matrix::new(self.layer_sizes[idx]+1, self.layer_sizes[idx + 1], weights)

    }

    fn compute_cost_grad(&self, data: &Matrix<f64>, outputs: &Matrix<f64>) -> Vec<f64> {
    	assert_eq!(data.cols(), self.layer_sizes[0]);

    	let mut forward_weights = Vec::with_capacity(self.layer_sizes.len()-1);
    	let mut activations = Vec::with_capacity(self.layer_sizes.len());

    	activations.push(data.clone());

    	// Forward propagation
    	{
	    	let mut a = self.get_layer_weights(0) * data;
	    	activations.push(a.clone());

	    	for l in 1..self.layer_sizes.len()-1 {
	    		let z = self.get_layer_weights(l) * a.clone();
	    		forward_weights.push(z.clone());

	    		activations.push(z.apply(&sigmoid));
	    	}

	    }

	    let mut deltas = Vec::with_capacity(self.layer_sizes.len()-1);
	    // Backward propagation
	    {
	    	let mut delta = &activations[self.layer_sizes.len()-1] - outputs;
	    	deltas.push(delta.clone());

	    	for l in (0..self.layer_sizes.len()-1).rev() {
	    		let g = forward_weights[l].clone().apply(&sigmoid_grad);
	    		delta = (self.get_layer_weights(l).transpose() * delta.clone()).elemul(&g);
	    		deltas.push(delta.clone());
	    	}
	    }

	    let mut grad = Vec::with_capacity(self.layer_sizes.len()-1);
	   	let mut capacity = 0;

	   	for l in 0..self.layer_sizes.len()-1 {
	   		// Probably need to remove the first column on the deltas.
	   		let g = deltas[l].clone() * activations[l].clone().transpose();
	   		capacity += g.cols() * g.rows();
	   		grad.push(g / (data.rows() as f64));
	   	}

	   	let mut gradients = Vec::with_capacity(capacity);

	   	for g in grad {
	   		gradients.append(&mut g.data.clone());
	   	}

	   	gradients
    }

}

impl<'a> SupModel<Matrix<f64>, Vector<usize>> for NeuralNet<'a> {
    fn predict(&self, data: Matrix<f64>) -> Vector<usize> {
        unimplemented!()
    }

    fn train(&mut self, data: Matrix<f64>, values: Vector<usize>) {
        unimplemented!()
    }
}

/// Sigmoid function.
///
/// Returns 1 / ( 1 + e^-t).
fn sigmoid(t: f64) -> f64 {
	1.0 / (1.0 + (-t).exp() )
}

/// Gradient of sigmoid function.
///
/// Evaluates to (1 - e^-t) / (1 + e^-t)^2
fn sigmoid_grad(t: f64) -> f64 {
	sigmoid(t) * (1f64-sigmoid(t))
}<|endoftext|>


<commit_before>/*
 * Copyright (C) 2017-2018 Ryan Huang
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published
 * by the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

mod board;
mod board_parts;
mod sys;
mod vec2;

use imgui::{self, ImGuiCond, ImStr, ImVec2, Ui};

use self::board::board;
pub use self::sys::run;
use self::vec2::Vec2;
use model::{Color, ColorMap, Model, Player};
use update::Event;

pub fn draw(ui: &Ui, size: (f32, f32), model: &Model) -> Option<Event> {
    unsafe {
        imgui::sys::igPushStyleVar(imgui::sys::ImGuiStyleVar::WindowRounding, 0.0);
    }

    let mut event = None;

    ui.main_menu_bar(|| {
        ui.menu(im_str!("Game")).build(|| {
            ui.menu_item(im_str!("New game")).enabled(false).build();

            use self::Player::*;
            if ui.menu_item(im_str!("Human vs. Human")).build() {
                insert_if_empty(&mut event, Event::NewGame(ColorMap::new(Human, Human)));
            }
            if ui.menu_item(im_str!("Human vs. Computer")).build() {
                insert_if_empty(&mut event, Event::NewGame(ColorMap::new(Human, Computer)));
            }
            if ui.menu_item(im_str!("Computer vs. Human")).build() {
                insert_if_empty(&mut event, Event::NewGame(ColorMap::new(Computer, Human)));
            }
            if ui.menu_item(im_str!("Computer vs. Computer")).build() {
                insert_if_empty(
                    &mut event,
                    Event::NewGame(ColorMap::new(Computer, Computer)),
                );
            }
            ui.separator();
            if ui.menu_item(im_str!("Quit")).build() {
                insert_if_empty(&mut event, Event::Quit);
            }
        });

        let mut depth = model.ai_search_depth as i32;
        ui.menu(im_str!("Computer")).build(|| {
            ui.slider_int(im_str!("Search depth"), &mut depth, 1, 7)
                .build();
            if ui.is_item_hovered() {
                ui.tooltip_text("How many moves ahead the computer will search.\nFewer moves is faster and easier, while more moves is slower and more difficult.");
            }
        });
        if depth as u8 != model.ai_search_depth {
            insert_if_empty(&mut event, Event::SetAIDepth(depth as u8));
        }
    });

    ui.window(im_str!("Coerceo"))
        .size(size, ImGuiCond::Always)
        .position((0.0, 19.0), ImGuiCond::Once)
        .title_bar(false)
        .resizable(false)
        .movable(false)
        .build(|| {
            ui.text("Welcome to Coerceo!");
            ui.text(format!(
                "{:?} vs. {:?}",
                model.players.white, model.players.black
            ));

            if let Some(click) = board(ui, model, Vec2::new(size.0 - 16.0, size.1 - 170.0)) {
                insert_if_empty(&mut event, click);
            }

            let display_vitals = || {
                ui.text(format!(
                    "{:?} has {} piece(s) left and {} captured hex(es).",
                    Color::White,
                    model.board.pieces(Color::White),
                    model.board.hexes(Color::White),
                ));
                ui.text(format!(
                    "{:?} has {} piece(s) left and {} captured hex(es).",
                    Color::Black,
                    model.board.pieces(Color::Black),
                    model.board.hexes(Color::Black),
                ));
            };

            let button_size = Vec2::new(120.0, 20.0);
            use model::Outcome::*;
            match model.outcome {
                Win(color) => {
                    ui.text(format!("{:?} wins!", color));
                    display_vitals();
                    if model.can_undo() && ui.button(im_str!("Undo"), button_size) {
                        insert_if_empty(&mut event, Event::Undo);
                    }
                }
                InProgress => {
                    ui.text(format!(
                        "It's {:?}'s turn. ({})",
                        model.board.turn,
                        match model.current_player() {
                            Player::Computer => "Computer",
                            Player::Human => "Human",
                        }
                    ));
                    display_vitals();

                    horz_button_layout(
                        ui,
                        vec![
                            (model.can_undo(), im_str!("Undo"), Event::Undo),
                            (model.can_redo(), im_str!("Redo"), Event::Redo),
                        ],
                        &button_size,
                        &mut event,
                    );
                    let is_human_player = model.current_player() == Player::Human;
                    horz_button_layout(
                        ui,
                        vec![
                            (is_human_player, im_str!("Resign"), Event::Resign),
                            (
                                model.board.can_exchange() && is_human_player,
                                if model.exchanging {
                                    im_str!("Stop Exchanging")
                                } else {
                                    im_str!("Exchange")
                                },
                                Event::Exchange,
                            ),
                        ],
                        &button_size,
                        &mut event,
                    );
                }
                // Draw cases
                _ => {
                    let message = match model.outcome {
                        DrawStalemate => "It's a draw by stalemate!",
                        DrawThreefoldRepetition => "It's a draw by threefold repetition!",
                        DrawInsufficientMaterial => "It's a draw by insufficient material!",
                        _ => unreachable!(),
                    };
                    ui.text(message);
                    display_vitals();
                    if model.can_undo() && ui.button(im_str!("Undo"), button_size) {
                        insert_if_empty(&mut event, Event::Undo);
                    }
                }
            }
        });

    unsafe {
        imgui::sys::igPopStyleVar(1);
    }

    event
}

fn horz_button_layout(
    ui: &Ui,
    buttons: Vec<(bool, &ImStr, Event)>,
    size: &Vec2,
    event: &mut Option<Event>,
) {
    if !buttons.iter().any(|&(show, _, _)| show) {
        return;
    }

    let size: ImVec2 = (*size).into();

    for (show, label, action) in buttons {
        if show {
            if ui.button(label, size) {
                insert_if_empty(event, action);
            }
        } else {
            ui.dummy(size);
        }
        ui.same_line(0.0);
    }
    ui.new_line();
}

fn insert_if_empty<T>(a: &mut Option<T>, b: T) {
    if a.is_none() {
        *a = Some(b);
    }
}
<commit_msg>Use lossless i32 conversion<commit_after>/*
 * Copyright (C) 2017-2018 Ryan Huang
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published
 * by the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

mod board;
mod board_parts;
mod sys;
mod vec2;

use imgui::{self, ImGuiCond, ImStr, ImVec2, Ui};

use self::board::board;
pub use self::sys::run;
use self::vec2::Vec2;
use model::{Color, ColorMap, Model, Player};
use update::Event;

pub fn draw(ui: &Ui, size: (f32, f32), model: &Model) -> Option<Event> {
    unsafe {
        imgui::sys::igPushStyleVar(imgui::sys::ImGuiStyleVar::WindowRounding, 0.0);
    }

    let mut event = None;

    ui.main_menu_bar(|| {
        ui.menu(im_str!("Game")).build(|| {
            ui.menu_item(im_str!("New game")).enabled(false).build();

            use self::Player::*;
            if ui.menu_item(im_str!("Human vs. Human")).build() {
                insert_if_empty(&mut event, Event::NewGame(ColorMap::new(Human, Human)));
            }
            if ui.menu_item(im_str!("Human vs. Computer")).build() {
                insert_if_empty(&mut event, Event::NewGame(ColorMap::new(Human, Computer)));
            }
            if ui.menu_item(im_str!("Computer vs. Human")).build() {
                insert_if_empty(&mut event, Event::NewGame(ColorMap::new(Computer, Human)));
            }
            if ui.menu_item(im_str!("Computer vs. Computer")).build() {
                insert_if_empty(
                    &mut event,
                    Event::NewGame(ColorMap::new(Computer, Computer)),
                );
            }
            ui.separator();
            if ui.menu_item(im_str!("Quit")).build() {
                insert_if_empty(&mut event, Event::Quit);
            }
        });

        let mut depth = i32::from(model.ai_search_depth);
        ui.menu(im_str!("Computer")).build(|| {
            ui.slider_int(im_str!("Search depth"), &mut depth, 1, 7)
                .build();
            if ui.is_item_hovered() {
                ui.tooltip_text("How many moves ahead the computer will search.\nFewer moves is faster and easier, while more moves is slower and more difficult.");
            }
        });
        if depth as u8 != model.ai_search_depth {
            insert_if_empty(&mut event, Event::SetAIDepth(depth as u8));
        }
    });

    ui.window(im_str!("Coerceo"))
        .size(size, ImGuiCond::Always)
        .position((0.0, 19.0), ImGuiCond::Once)
        .title_bar(false)
        .resizable(false)
        .movable(false)
        .build(|| {
            ui.text("Welcome to Coerceo!");
            ui.text(format!(
                "{:?} vs. {:?}",
                model.players.white, model.players.black
            ));

            if let Some(click) = board(ui, model, Vec2::new(size.0 - 16.0, size.1 - 170.0)) {
                insert_if_empty(&mut event, click);
            }

            let display_vitals = || {
                ui.text(format!(
                    "{:?} has {} piece(s) left and {} captured hex(es).",
                    Color::White,
                    model.board.pieces(Color::White),
                    model.board.hexes(Color::White),
                ));
                ui.text(format!(
                    "{:?} has {} piece(s) left and {} captured hex(es).",
                    Color::Black,
                    model.board.pieces(Color::Black),
                    model.board.hexes(Color::Black),
                ));
            };

            let button_size = Vec2::new(120.0, 20.0);
            use model::Outcome::*;
            match model.outcome {
                Win(color) => {
                    ui.text(format!("{:?} wins!", color));
                    display_vitals();
                    if model.can_undo() && ui.button(im_str!("Undo"), button_size) {
                        insert_if_empty(&mut event, Event::Undo);
                    }
                }
                InProgress => {
                    ui.text(format!(
                        "It's {:?}'s turn. ({})",
                        model.board.turn,
                        match model.current_player() {
                            Player::Computer => "Computer",
                            Player::Human => "Human",
                        }
                    ));
                    display_vitals();

                    horz_button_layout(
                        ui,
                        vec![
                            (model.can_undo(), im_str!("Undo"), Event::Undo),
                            (model.can_redo(), im_str!("Redo"), Event::Redo),
                        ],
                        &button_size,
                        &mut event,
                    );
                    let is_human_player = model.current_player() == Player::Human;
                    horz_button_layout(
                        ui,
                        vec![
                            (is_human_player, im_str!("Resign"), Event::Resign),
                            (
                                model.board.can_exchange() && is_human_player,
                                if model.exchanging {
                                    im_str!("Stop Exchanging")
                                } else {
                                    im_str!("Exchange")
                                },
                                Event::Exchange,
                            ),
                        ],
                        &button_size,
                        &mut event,
                    );
                }
                // Draw cases
                _ => {
                    let message = match model.outcome {
                        DrawStalemate => "It's a draw by stalemate!",
                        DrawThreefoldRepetition => "It's a draw by threefold repetition!",
                        DrawInsufficientMaterial => "It's a draw by insufficient material!",
                        _ => unreachable!(),
                    };
                    ui.text(message);
                    display_vitals();
                    if model.can_undo() && ui.button(im_str!("Undo"), button_size) {
                        insert_if_empty(&mut event, Event::Undo);
                    }
                }
            }
        });

    unsafe {
        imgui::sys::igPopStyleVar(1);
    }

    event
}

fn horz_button_layout(
    ui: &Ui,
    buttons: Vec<(bool, &ImStr, Event)>,
    size: &Vec2,
    event: &mut Option<Event>,
) {
    if !buttons.iter().any(|&(show, _, _)| show) {
        return;
    }

    let size: ImVec2 = (*size).into();

    for (show, label, action) in buttons {
        if show {
            if ui.button(label, size) {
                insert_if_empty(event, action);
            }
        } else {
            ui.dummy(size);
        }
        ui.same_line(0.0);
    }
    ui.new_line();
}

fn insert_if_empty<T>(a: &mut Option<T>, b: T) {
    if a.is_none() {
        *a = Some(b);
    }
}
<|endoftext|>


<commit_before>// A port of the simplistic benchmark from
//
//    http://github.com/PaulKeeble/ScalaVErlangAgents
//
// I *think* it's the same, more or less.

#![crate_name = "msgsend"]
#![allow(unstable)]

extern crate time;
extern crate zmq;

use std::os;
use std::thread::Thread;
use std::sync::mpsc::{channel, Sender, Receiver};

fn server(mut pull_socket: zmq::Socket, mut push_socket: zmq::Socket, mut workers: u64) {
    let mut count = 0;
    let mut msg = zmq::Message::new().unwrap();

    while workers != 0 {
        match pull_socket.recv(&mut msg, 0) {
            Err(e) => panic!(e),
            Ok(()) => {
                let s = msg.as_str().unwrap();
                if s.is_empty() {
                    workers -= 1;
                } else {
                    count += s.parse().unwrap();
                }
            }
        }
    }

    match push_socket.send_str(count.to_string().as_slice(), 0) {
        Ok(()) => { }
        Err(e) => panic!(e),
    }
}

fn spawn_server(ctx: &mut zmq::Context, workers: u64) -> Sender<()> {
    let mut pull_socket = ctx.socket(zmq::PULL).unwrap();
    let mut push_socket = ctx.socket(zmq::PUSH).unwrap();

    pull_socket.bind("inproc://server-pull").unwrap();
    push_socket.bind("inproc://server-push").unwrap();

    // Spawn the server.
    let (ready_tx, ready_rx) = channel();
    let (start_tx, start_rx) = channel();

    Thread::spawn(move|| {
        // Let the main thread know we're ready.
        ready_tx.send(()).unwrap();

        // Wait until we need to start.
        start_rx.recv().unwrap();

        server(pull_socket, push_socket, workers);
    });

    // Wait for the server to start.
    ready_rx.recv().unwrap();

    start_tx
}

fn worker(mut push_socket: zmq::Socket, count: u64) {
    for _ in range(0, count) {
        push_socket.send_str(100.to_string().as_slice(), 0).unwrap();
    }

    // Let the server know we're done.
    push_socket.send_str("", 0).unwrap();
}

fn spawn_worker(ctx: &mut zmq::Context, count: u64) -> Receiver<()> {
    let mut push_socket = ctx.socket(zmq::PUSH).unwrap();

    push_socket.connect("inproc://server-pull").unwrap();
    //push_socket.connect("tcp://127.0.0.1:3456").unwrap();

    // Spawn the worker.
    let (tx, rx) = channel();
    Thread::spawn(move|| {
        // Let the main thread we're ready.
        tx.send(()).unwrap();

        worker(push_socket, count);

        tx.send(()).unwrap();
    });

    // Wait for the worker to start.
    rx.recv().unwrap();

    rx
}

fn run(ctx: &mut zmq::Context, size: u64, workers: u64) {
    let start_ch = spawn_server(ctx, workers);

    // Create some command/control sockets.
    let mut push_socket = ctx.socket(zmq::PUSH).unwrap();
    let mut pull_socket = ctx.socket(zmq::PULL).unwrap();

    push_socket.connect("inproc://server-pull").unwrap();
    pull_socket.connect("inproc://server-push").unwrap();
    //push_socket.connect("tcp://127.0.0.1:3456").unwrap();
    //pull_socket.connect("tcp://127.0.0.1:3457").unwrap();

    // Spawn all the workers.
    let mut worker_results = Vec::new();
    for _ in range(0, workers) {
        worker_results.push(spawn_worker(ctx, size / workers));
    }

    let start = time::precise_time_s();

    start_ch.send(()).unwrap();

    // Block until all the workers finish.
    for po in worker_results.iter() {
        po.recv().unwrap();
    }

    // Receive the final count.
    let result: i32 = match pull_socket.recv_msg(0) {
        Ok(msg) => {
            let msg_str: &str = msg.as_str().unwrap();
            msg_str.parse().unwrap()
        },
        Err(e) => panic!(e),
    };

    let end = time::precise_time_s();
    let elapsed = end - start;

    println!("Count is {}", result);
    println!("Test took {} seconds", elapsed);
    let thruput = ((size / workers * workers) as f64) / elapsed;
    println!("Throughput={} per sec", thruput);
}

fn main() {
    let args = os::args();

    let args = if os::getenv("RUST_BENCH").is_some() {
        vec!("".to_string(), "1000000".to_string(), "10000".to_string())
    } else if args.len() <= 1 {
        vec!("".to_string(), "10000".to_string(), "4".to_string())
    } else {
        args
    };

    let size = args[1].as_slice().parse().unwrap();
    let workers = args[2].as_slice().parse().unwrap();

    let mut ctx = zmq::Context::new();

    run(&mut ctx, size, workers);
}
<commit_msg>Thread -> thread<commit_after>// A port of the simplistic benchmark from
//
//    http://github.com/PaulKeeble/ScalaVErlangAgents
//
// I *think* it's the same, more or less.

#![crate_name = "msgsend"]
#![allow(unstable)]

extern crate time;
extern crate zmq;

use std::os;
use std::thread;
use std::sync::mpsc::{channel, Sender, Receiver};

fn server(mut pull_socket: zmq::Socket, mut push_socket: zmq::Socket, mut workers: u64) {
    let mut count = 0;
    let mut msg = zmq::Message::new().unwrap();

    while workers != 0 {
        match pull_socket.recv(&mut msg, 0) {
            Err(e) => panic!(e),
            Ok(()) => {
                let s = msg.as_str().unwrap();
                if s.is_empty() {
                    workers -= 1;
                } else {
                    count += s.parse().unwrap();
                }
            }
        }
    }

    match push_socket.send_str(count.to_string().as_slice(), 0) {
        Ok(()) => { }
        Err(e) => panic!(e),
    }
}

fn spawn_server(ctx: &mut zmq::Context, workers: u64) -> Sender<()> {
    let mut pull_socket = ctx.socket(zmq::PULL).unwrap();
    let mut push_socket = ctx.socket(zmq::PUSH).unwrap();

    pull_socket.bind("inproc://server-pull").unwrap();
    push_socket.bind("inproc://server-push").unwrap();

    // Spawn the server.
    let (ready_tx, ready_rx) = channel();
    let (start_tx, start_rx) = channel();

    thread::spawn(move|| {
        // Let the main thread know we're ready.
        ready_tx.send(()).unwrap();

        // Wait until we need to start.
        start_rx.recv().unwrap();

        server(pull_socket, push_socket, workers);
    });

    // Wait for the server to start.
    ready_rx.recv().unwrap();

    start_tx
}

fn worker(mut push_socket: zmq::Socket, count: u64) {
    for _ in range(0, count) {
        push_socket.send_str(100.to_string().as_slice(), 0).unwrap();
    }

    // Let the server know we're done.
    push_socket.send_str("", 0).unwrap();
}

fn spawn_worker(ctx: &mut zmq::Context, count: u64) -> Receiver<()> {
    let mut push_socket = ctx.socket(zmq::PUSH).unwrap();

    push_socket.connect("inproc://server-pull").unwrap();
    //push_socket.connect("tcp://127.0.0.1:3456").unwrap();

    // Spawn the worker.
    let (tx, rx) = channel();
    thread::spawn(move|| {
        // Let the main thread we're ready.
        tx.send(()).unwrap();

        worker(push_socket, count);

        tx.send(()).unwrap();
    });

    // Wait for the worker to start.
    rx.recv().unwrap();

    rx
}

fn run(ctx: &mut zmq::Context, size: u64, workers: u64) {
    let start_ch = spawn_server(ctx, workers);

    // Create some command/control sockets.
    let mut push_socket = ctx.socket(zmq::PUSH).unwrap();
    let mut pull_socket = ctx.socket(zmq::PULL).unwrap();

    push_socket.connect("inproc://server-pull").unwrap();
    pull_socket.connect("inproc://server-push").unwrap();
    //push_socket.connect("tcp://127.0.0.1:3456").unwrap();
    //pull_socket.connect("tcp://127.0.0.1:3457").unwrap();

    // Spawn all the workers.
    let mut worker_results = Vec::new();
    for _ in range(0, workers) {
        worker_results.push(spawn_worker(ctx, size / workers));
    }

    let start = time::precise_time_s();

    start_ch.send(()).unwrap();

    // Block until all the workers finish.
    for po in worker_results.iter() {
        po.recv().unwrap();
    }

    // Receive the final count.
    let result: i32 = match pull_socket.recv_msg(0) {
        Ok(msg) => {
            let msg_str: &str = msg.as_str().unwrap();
            msg_str.parse().unwrap()
        },
        Err(e) => panic!(e),
    };

    let end = time::precise_time_s();
    let elapsed = end - start;

    println!("Count is {}", result);
    println!("Test took {} seconds", elapsed);
    let thruput = ((size / workers * workers) as f64) / elapsed;
    println!("Throughput={} per sec", thruput);
}

fn main() {
    let args = os::args();

    let args = if os::getenv("RUST_BENCH").is_some() {
        vec!("".to_string(), "1000000".to_string(), "10000".to_string())
    } else if args.len() <= 1 {
        vec!("".to_string(), "10000".to_string(), "4".to_string())
    } else {
        args
    };

    let size = args[1].as_slice().parse().unwrap();
    let workers = args[2].as_slice().parse().unwrap();

    let mut ctx = zmq::Context::new();

    run(&mut ctx, size, workers);
}
<|endoftext|>


<commit_before>// Copyright 2016 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

extern crate byteorder;
extern crate cgmath;
extern crate clap;
extern crate collision;
extern crate pbr;
extern crate point_viewer;
extern crate protobuf;
extern crate scoped_pool;

use cgmath::{EuclideanSpace, Point3};
use collision::{Aabb, Aabb3};
use pbr::ProgressBar;
use point_viewer::{InternalIterator, Point};
use point_viewer::errors::*;
use point_viewer::math::Cube;
use point_viewer::octree;
use point_viewer::ply::PlyIterator;
use point_viewer::proto;
use point_viewer::pts::PtsIterator;
use protobuf::Message;
use scoped_pool::{Pool, Scope};
use std::collections::{HashMap, HashSet};
use std::fs::{self, File};
use std::io::{BufWriter, Stdout};
use std::path::PathBuf;
use std::sync::mpsc;

const UPDATE_COUNT: i64 = 100000;
const MAX_POINTS_PER_NODE: i64 = 100000;

// Return a list a leaf nodes and a list of nodes to be splitted further.
fn split<P>(
    octree_meta: &octree::OctreeMeta,
    node_id: &octree::NodeId,
    stream: P,
) -> (Vec<octree::NodeId>, Vec<octree::NodeId>)
where
    P: InternalIterator,
{
    let mut children: Vec<Option<octree::NodeWriter>> =
        vec![None, None, None, None, None, None, None, None];
    match stream.size_hint() {
        Some(size) => println!(
            "Splitting {} which has {} points ({:.2}x MAX_POINTS_PER_NODE).",
            node_id,
            size,
            size as f64 / MAX_POINTS_PER_NODE as f64
        ),
        None => println!(
            "Splitting {} which has an unknown number of points.",
            node_id
        ),
    };

    let bounding_cube = node_id.find_bounding_cube(&Cube::bounding(&octree_meta.bounding_box));
    stream.for_each(|p| {
        let child_index = octree::ChildIndex::from_bounding_cube(&bounding_cube, &p.position);
        let array_index = child_index.as_u8() as usize;
        if children[array_index].is_none() {
            children[array_index] = Some(octree::NodeWriter::new(
                &octree_meta,
                &node_id.get_child_id(child_index),
            ));
        }
        children[array_index].as_mut().unwrap().write(p);
    });

    // Remove the node file on disk by reopening the node and immediately dropping it again without
    // writing a point. This only saves some disk space during processing - all nodes will be
    // rewritten by subsampling the children in the second step anyways. We also ignore file
    // removing error. For example, we never write out the root, so it cannot be removed.
    octree::NodeWriter::new(&octree_meta, &node_id);

    let mut leaf_nodes = Vec::new();
    let mut split_nodes = Vec::new();
    for (child_index, c) in children.into_iter().enumerate() {
        if c.is_none() {
            continue;
        }
        let c = c.unwrap();
        let child_id = node_id.get_child_id(octree::ChildIndex::from_u8(child_index as u8));

        if should_split_node(&child_id, c.num_written(), octree_meta) {
            split_nodes.push(child_id);
        } else {
            leaf_nodes.push(child_id);
        }
    }
    (leaf_nodes, split_nodes)
}

fn should_split_node(
    id: &octree::NodeId,
    num_points: i64,
    octree_meta: &octree::OctreeMeta,
) -> bool {
    if num_points <= MAX_POINTS_PER_NODE {
        return false;
    }
    let bounding_cube = id.find_bounding_cube(&Cube::bounding(&octree_meta.bounding_box));
    if f64::from(bounding_cube.edge_length()) <= octree_meta.resolution {
        // TODO(hrapp): If the data has billion of points in this small spot, performance will
        // greatly suffer if we display it. Drop points?
        println!(
            "Node {} which has {} points ({:.2}x MAX_POINTS_PER_NODE) \
             is too small to be split, keeping all points.",
            id,
            num_points,
            num_points as f64 / MAX_POINTS_PER_NODE as f64
        );
        return false;
    }
    true
}

fn split_node<'a, 'b: 'a, P>(
    scope: &Scope<'a>,
    octree_meta: &'a octree::OctreeMeta,
    node_id: &octree::NodeId,
    stream: P,
    leaf_nodes_sender: &mpsc::Sender<octree::NodeId>,
) where
    P: InternalIterator,
{
    let (leaf_nodes, split_nodes) = split(octree_meta, &node_id, stream);
    for child_id in split_nodes {
        let leaf_nodes_sender_clone = leaf_nodes_sender.clone();
        scope.recurse(move |scope| {
            let stream = octree::NodeIterator::from_disk(&octree_meta, &child_id).unwrap();
            split_node(
                scope,
                octree_meta,
                &child_id,
                stream,
                &leaf_nodes_sender_clone,
            );
        });
    }

    for id in leaf_nodes {
        leaf_nodes_sender.send(id).unwrap();
    }
}

fn subsample_children_into(
    octree_meta: &octree::OctreeMeta,
    node_id: &octree::NodeId,
    nodes_sender: &mpsc::Sender<(octree::NodeId, i64)>,
) -> Result<()> {
    let mut parent_writer = octree::NodeWriter::new(&octree_meta, &node_id);
    println!("Creating {} from subsampling children.", node_id);
    for i in 0..8 {
        let child_id = node_id.get_child_id(octree::ChildIndex::from_u8(i));
        let node_iterator = match octree::NodeIterator::from_disk(&octree_meta, &child_id) {
            Ok(node_iterator) => node_iterator,
            Err(Error(ErrorKind::NodeNotFound, _)) => continue,
            Err(err) => return Err(err),
        };

        // We read all points into memory, because the new node writer will rewrite this child's
        // file(s).
        let mut points = Vec::with_capacity(node_iterator.size_hint().unwrap());
        node_iterator.for_each(|p| points.push((*p).clone()));

        let mut child_writer = octree::NodeWriter::new(&octree_meta, &child_id);
        for (idx, p) in points.into_iter().enumerate() {
            if idx % 8 == 0 {
                parent_writer.write(&p);
            } else {
                child_writer.write(&p);
            }
        }
        // Update child.
        nodes_sender
            .send((child_id, child_writer.num_written()))
            .unwrap();
    }
    // Make sure the root node is also tracked as an existing node.
    if node_id.level() == 0 {
        nodes_sender
            .send((*node_id, parent_writer.num_written()))
            .unwrap();
    }
    Ok(())
}

#[derive(Debug)]
enum InputFile {
    Ply(PathBuf),
    Pts(PathBuf),
}

enum InputFileIterator {
    Ply(PlyIterator),
    Pts(PtsIterator),
}

impl InternalIterator for InputFileIterator {
    fn size_hint(&self) -> Option<usize> {
        match *self {
            InputFileIterator::Ply(ref p) => p.size_hint(),
            InputFileIterator::Pts(ref p) => p.size_hint(),
        }
    }

    fn for_each<F: FnMut(&Point)>(self, f: F) {
        match self {
            InputFileIterator::Ply(p) => p.for_each(f),
            InputFileIterator::Pts(p) => p.for_each(f),
        }
    }
}

fn make_stream(input: &InputFile) -> (InputFileIterator, Option<pbr::ProgressBar<Stdout>>) {
    let stream = match *input {
        InputFile::Ply(ref filename) => InputFileIterator::Ply(PlyIterator::new(filename).unwrap()),
        InputFile::Pts(ref filename) => InputFileIterator::Pts(PtsIterator::new(filename)),
    };

    let progress_bar = match stream.size_hint() {
        Some(size) => Some(ProgressBar::new(size as u64)),
        None => None,
    };
    (stream, progress_bar)
}

/// Returns the bounding_box and the number of the points in 'input'.
fn find_bounding_box(input: &InputFile) -> Aabb3<f32> {
    let mut num_points = 0i64;
    let mut bounding_box = Aabb3::zero();
    let (stream, mut progress_bar) = make_stream(input);
    progress_bar
        .as_mut()
        .map(|pb| pb.message("Determining bounding box: "));

    stream.for_each(|p: &Point| {
        bounding_box = bounding_box.grow(Point3::from_vec(p.position));
        num_points += 1;
        if num_points % UPDATE_COUNT == 0 {
            progress_bar.as_mut().map(|pb| pb.add(UPDATE_COUNT as u64));
        }
    });
    progress_bar.map(|mut f| f.finish());
    bounding_box
}

fn main() {
    let matches = clap::App::new("build_octree")
        .args(&[
            clap::Arg::with_name("output_directory")
                .help("Output directory to write the octree into.")
                .long("output_directory")
                .required(true)
                .takes_value(true),
            clap::Arg::with_name("resolution")
                .help(
                    "Minimal precision that this point cloud should have. This decides \
                     on the number of bits used to encode each node.",
                )
                .long("resolution")
                .default_value("0.001"),
            clap::Arg::with_name("input")
                .help("PLY/PTS file to parse for the points.")
                .index(1)
                .required(true),
        ])
        .get_matches();

    let output_directory = &PathBuf::from(matches.value_of("output_directory").unwrap());
    let resolution = matches
        .value_of("resolution")
        .unwrap()
        .parse::<f64>()
        .expect("resolution could not be parsed as float.");

    let input = {
        let filename = PathBuf::from(matches.value_of("input").unwrap());
        match filename.extension().and_then(|s| s.to_str()) {
            Some("ply") => InputFile::Ply(filename.clone()),
            Some("pts") => InputFile::Pts(filename.clone()),
            other => panic!("Unknown input file format: {:?}", other),
        }
    };

    let bounding_box = find_bounding_box(&input);

    let octree_meta = &octree::OctreeMeta {
        bounding_box,
        resolution,
        directory: output_directory.clone(),
    };

    // Ignore errors, maybe directory is already there.
    let _ = fs::create_dir(output_directory);

    let mut meta = {
        let mut meta = proto::Meta::new();
        meta.mut_bounding_box()
            .mut_min()
            .set_x(bounding_box.min().x);
        meta.mut_bounding_box()
            .mut_min()
            .set_y(bounding_box.min().y);
        meta.mut_bounding_box()
            .mut_min()
            .set_z(bounding_box.min().z);
        meta.mut_bounding_box()
            .mut_max()
            .set_x(bounding_box.max().x);
        meta.mut_bounding_box()
            .mut_max()
            .set_y(bounding_box.max().y);
        meta.mut_bounding_box()
            .mut_max()
            .set_z(bounding_box.max().z);
        meta.set_resolution(resolution);
        meta.set_version(octree::CURRENT_VERSION);
        meta
    };

    println!("Creating octree structure.");
    let pool = Pool::new(10);

    let (leaf_nodes_sender, leaf_nodes_receiver) = mpsc::channel();
    pool.scoped(move |scope| {
        let (root_stream, _) = make_stream(&input);
        let root_node = octree::Node::root_with_bounding_cube(Cube::bounding(&bounding_box));
        split_node(
            scope,
            octree_meta,
            &root_node.id,
            root_stream,
            &leaf_nodes_sender,
        );
    });

    let mut nodes_to_subsample = Vec::new();
    let mut deepest_level = 0usize;
    for id in leaf_nodes_receiver {
        deepest_level = std::cmp::max(deepest_level, id.level());
        nodes_to_subsample.push(id);
    }
    let mut finished_nodes = HashMap::new();

    // sub sampling returns the list of finished nodes including all meta data
    // We start on the deepest level and work our way up the tree.
    for current_level in (0..deepest_level + 1).rev() {
        // All nodes on the same level can be subsampled in parallel.
        let res = nodes_to_subsample
            .into_iter()
            .partition(|n| n.level() == current_level);
        nodes_to_subsample = res.1;

        let mut parent_ids = HashSet::new();
        let mut subsample_nodes = Vec::new();
        for id in res.0 {
            let maybe_parent_id = id.parent_id();
            if maybe_parent_id.is_none() {
                // Only the root has no parents.
                assert_eq!(id.level(), 0);
                continue;
            }
            let parent_id = maybe_parent_id.unwrap();
            if parent_ids.contains(&parent_id) {
                continue;
            }
            parent_ids.insert(parent_id);
            subsample_nodes.push(parent_id);
        }

        let (finished_nodes_sender, finished_nodes_receiver) = mpsc::channel();
        pool.scoped(|scope| {
            for id in &subsample_nodes {
                let finished_nodes_sender_clone = finished_nodes_sender.clone();
                scope.execute(move || {
                    subsample_children_into(octree_meta, id, &finished_nodes_sender_clone).unwrap();
                });
            }
        });

        drop(finished_nodes_sender);
        for (id, num_points) in finished_nodes_receiver {
            if num_points > 0 {
                finished_nodes.insert(id, num_points);
            }
        }

        // The nodes that were just now created through sub-sampling will be required to create
        // their parents.
        nodes_to_subsample.extend(subsample_nodes.into_iter());
    }

    // Add all non-zero node meta data to meta.pb
    for (id, num_points) in finished_nodes {
        let bounding_cube = id.find_bounding_cube(&Cube::bounding(&octree_meta.bounding_box));
        let position_encoding =
            octree::PositionEncoding::new(&bounding_cube, octree_meta.resolution);

        let mut proto = proto::Node::new();
        proto.mut_id().set_level(id.level() as i32);
        proto.mut_id().set_index(id.index() as i64);
        proto.set_num_points(num_points);
        proto.set_position_encoding(position_encoding.to_proto());
        meta.mut_nodes().push(proto);
    }

    let mut buf_writer = BufWriter::new(File::create(&output_directory.join("meta.pb")).unwrap());
    meta.write_to_writer(&mut buf_writer).unwrap();
}
<commit_msg>clippy cleanup (#96)<commit_after>// Copyright 2016 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

extern crate byteorder;
extern crate cgmath;
extern crate clap;
extern crate collision;
extern crate pbr;
extern crate point_viewer;
extern crate protobuf;
extern crate scoped_pool;

use cgmath::{EuclideanSpace, Point3};
use collision::{Aabb, Aabb3};
use pbr::ProgressBar;
use point_viewer::{InternalIterator, Point};
use point_viewer::errors::*;
use point_viewer::math::Cube;
use point_viewer::octree;
use point_viewer::ply::PlyIterator;
use point_viewer::proto;
use point_viewer::pts::PtsIterator;
use protobuf::Message;
use scoped_pool::{Pool, Scope};
use std::collections::{HashMap, HashSet};
use std::fs::{self, File};
use std::io::{BufWriter, Stdout};
use std::path::PathBuf;
use std::sync::mpsc;

const UPDATE_COUNT: i64 = 100000;
const MAX_POINTS_PER_NODE: i64 = 100000;

// Return a list a leaf nodes and a list of nodes to be splitted further.
fn split<P>(
    octree_meta: &octree::OctreeMeta,
    node_id: &octree::NodeId,
    stream: P,
) -> (Vec<octree::NodeId>, Vec<octree::NodeId>)
where
    P: InternalIterator,
{
    let mut children: Vec<Option<octree::NodeWriter>> =
        vec![None, None, None, None, None, None, None, None];
    match stream.size_hint() {
        Some(size) => println!(
            "Splitting {} which has {} points ({:.2}x MAX_POINTS_PER_NODE).",
            node_id,
            size,
            size as f64 / MAX_POINTS_PER_NODE as f64
        ),
        None => println!(
            "Splitting {} which has an unknown number of points.",
            node_id
        ),
    };

    let bounding_cube = node_id.find_bounding_cube(&Cube::bounding(&octree_meta.bounding_box));
    stream.for_each(|p| {
        let child_index = octree::ChildIndex::from_bounding_cube(&bounding_cube, &p.position);
        let array_index = child_index.as_u8() as usize;
        if children[array_index].is_none() {
            children[array_index] = Some(octree::NodeWriter::new(
                octree_meta,
                &node_id.get_child_id(child_index),
            ));
        }
        children[array_index].as_mut().unwrap().write(p);
    });

    // Remove the node file on disk by reopening the node and immediately dropping it again without
    // writing a point. This only saves some disk space during processing - all nodes will be
    // rewritten by subsampling the children in the second step anyways. We also ignore file
    // removing error. For example, we never write out the root, so it cannot be removed.
    octree::NodeWriter::new(octree_meta, node_id);

    let mut leaf_nodes = Vec::new();
    let mut split_nodes = Vec::new();
    for (child_index, c) in children.into_iter().enumerate() {
        if c.is_none() {
            continue;
        }
        let c = c.unwrap();
        let child_id = node_id.get_child_id(octree::ChildIndex::from_u8(child_index as u8));

        if should_split_node(&child_id, c.num_written(), octree_meta) {
            split_nodes.push(child_id);
        } else {
            leaf_nodes.push(child_id);
        }
    }
    (leaf_nodes, split_nodes)
}

fn should_split_node(
    id: &octree::NodeId,
    num_points: i64,
    octree_meta: &octree::OctreeMeta,
) -> bool {
    if num_points <= MAX_POINTS_PER_NODE {
        return false;
    }
    let bounding_cube = id.find_bounding_cube(&Cube::bounding(&octree_meta.bounding_box));
    if f64::from(bounding_cube.edge_length()) <= octree_meta.resolution {
        // TODO(hrapp): If the data has billion of points in this small spot, performance will
        // greatly suffer if we display it. Drop points?
        println!(
            "Node {} which has {} points ({:.2}x MAX_POINTS_PER_NODE) \
             is too small to be split, keeping all points.",
            id,
            num_points,
            num_points as f64 / MAX_POINTS_PER_NODE as f64
        );
        return false;
    }
    true
}

fn split_node<'a, P>(
    scope: &Scope<'a>,
    octree_meta: &'a octree::OctreeMeta,
    node_id: &octree::NodeId,
    stream: P,
    leaf_nodes_sender: &mpsc::Sender<octree::NodeId>,
) where
    P: InternalIterator,
{
    let (leaf_nodes, split_nodes) = split(octree_meta, node_id, stream);
    for child_id in split_nodes {
        let leaf_nodes_sender_clone = leaf_nodes_sender.clone();
        scope.recurse(move |scope| {
            let stream = octree::NodeIterator::from_disk(octree_meta, &child_id).unwrap();
            split_node(
                scope,
                octree_meta,
                &child_id,
                stream,
                &leaf_nodes_sender_clone,
            );
        });
    }

    for id in leaf_nodes {
        leaf_nodes_sender.send(id).unwrap();
    }
}

fn subsample_children_into(
    octree_meta: &octree::OctreeMeta,
    node_id: &octree::NodeId,
    nodes_sender: &mpsc::Sender<(octree::NodeId, i64)>,
) -> Result<()> {
    let mut parent_writer = octree::NodeWriter::new(octree_meta, node_id);
    println!("Creating {} from subsampling children.", node_id);
    for i in 0..8 {
        let child_id = node_id.get_child_id(octree::ChildIndex::from_u8(i));
        let node_iterator = match octree::NodeIterator::from_disk(octree_meta, &child_id) {
            Ok(node_iterator) => node_iterator,
            Err(Error(ErrorKind::NodeNotFound, _)) => continue,
            Err(err) => return Err(err),
        };

        // We read all points into memory, because the new node writer will rewrite this child's
        // file(s).
        let mut points = Vec::with_capacity(node_iterator.size_hint().unwrap());
        node_iterator.for_each(|p| points.push((*p).clone()));

        let mut child_writer = octree::NodeWriter::new(octree_meta, &child_id);
        for (idx, p) in points.into_iter().enumerate() {
            if idx % 8 == 0 {
                parent_writer.write(&p);
            } else {
                child_writer.write(&p);
            }
        }
        // Update child.
        nodes_sender
            .send((child_id, child_writer.num_written()))
            .unwrap();
    }
    // Make sure the root node is also tracked as an existing node.
    if node_id.level() == 0 {
        nodes_sender
            .send((*node_id, parent_writer.num_written()))
            .unwrap();
    }
    Ok(())
}

#[derive(Debug)]
enum InputFile {
    Ply(PathBuf),
    Pts(PathBuf),
}

enum InputFileIterator {
    Ply(PlyIterator),
    Pts(PtsIterator),
}

impl InternalIterator for InputFileIterator {
    fn size_hint(&self) -> Option<usize> {
        match *self {
            InputFileIterator::Ply(ref p) => p.size_hint(),
            InputFileIterator::Pts(ref p) => p.size_hint(),
        }
    }

    fn for_each<F: FnMut(&Point)>(self, f: F) {
        match self {
            InputFileIterator::Ply(p) => p.for_each(f),
            InputFileIterator::Pts(p) => p.for_each(f),
        }
    }
}

fn make_stream(input: &InputFile) -> (InputFileIterator, Option<pbr::ProgressBar<Stdout>>) {
    let stream = match *input {
        InputFile::Ply(ref filename) => InputFileIterator::Ply(PlyIterator::new(filename).unwrap()),
        InputFile::Pts(ref filename) => InputFileIterator::Pts(PtsIterator::new(filename)),
    };

    let progress_bar = match stream.size_hint() {
        Some(size) => Some(ProgressBar::new(size as u64)),
        None => None,
    };
    (stream, progress_bar)
}

/// Returns the bounding box containing all points
fn find_bounding_box(input: &InputFile) -> Aabb3<f32> {
    let mut num_points = 0i64;
    let mut bounding_box = Aabb3::zero();
    let (stream, mut progress_bar) = make_stream(input);
    progress_bar
        .as_mut()
        .map(|pb| pb.message("Determining bounding box: "));

    stream.for_each(|p: &Point| {
        bounding_box = bounding_box.grow(Point3::from_vec(p.position));
        num_points += 1;
        if num_points % UPDATE_COUNT == 0 {
            progress_bar.as_mut().map(|pb| pb.add(UPDATE_COUNT as u64));
        }
    });
    progress_bar.map(|mut f| f.finish());
    bounding_box
}

fn main() {
    let matches = clap::App::new("build_octree")
        .args(&[
            clap::Arg::with_name("output_directory")
                .help("Output directory to write the octree into.")
                .long("output_directory")
                .required(true)
                .takes_value(true),
            clap::Arg::with_name("resolution")
                .help(
                    "Minimal precision that this point cloud should have. This decides \
                     on the number of bits used to encode each node.",
                )
                .long("resolution")
                .default_value("0.001"),
            clap::Arg::with_name("input")
                .help("PLY/PTS file to parse for the points.")
                .index(1)
                .required(true),
        ])
        .get_matches();

    let output_directory = &PathBuf::from(matches.value_of("output_directory").unwrap());
    let resolution = matches
        .value_of("resolution")
        .unwrap()
        .parse::<f64>()
        .expect("resolution could not be parsed as float.");

    let input = {
        let filename = PathBuf::from(matches.value_of("input").unwrap());
        match filename.extension().and_then(|s| s.to_str()) {
            Some("ply") => InputFile::Ply(filename.clone()),
            Some("pts") => InputFile::Pts(filename.clone()),
            other => panic!("Unknown input file format: {:?}", other),
        }
    };

    let bounding_box = find_bounding_box(&input);

    let octree_meta = &octree::OctreeMeta {
        bounding_box,
        resolution,
        directory: output_directory.clone(),
    };

    // Ignore errors, maybe directory is already there.
    let _ = fs::create_dir(output_directory);

    let mut meta = {
        let mut meta = proto::Meta::new();
        meta.mut_bounding_box()
            .mut_min()
            .set_x(bounding_box.min().x);
        meta.mut_bounding_box()
            .mut_min()
            .set_y(bounding_box.min().y);
        meta.mut_bounding_box()
            .mut_min()
            .set_z(bounding_box.min().z);
        meta.mut_bounding_box()
            .mut_max()
            .set_x(bounding_box.max().x);
        meta.mut_bounding_box()
            .mut_max()
            .set_y(bounding_box.max().y);
        meta.mut_bounding_box()
            .mut_max()
            .set_z(bounding_box.max().z);
        meta.set_resolution(resolution);
        meta.set_version(octree::CURRENT_VERSION);
        meta
    };

    println!("Creating octree structure.");
    let pool = Pool::new(10);

    let (leaf_nodes_sender, leaf_nodes_receiver) = mpsc::channel();
    pool.scoped(move |scope| {
        let (root_stream, _) = make_stream(&input);
        let root_node = octree::Node::root_with_bounding_cube(Cube::bounding(&bounding_box));
        split_node(
            scope,
            octree_meta,
            &root_node.id,
            root_stream,
            &leaf_nodes_sender,
        );
    });

    let mut nodes_to_subsample = Vec::new();
    let mut deepest_level = 0usize;
    for id in leaf_nodes_receiver {
        deepest_level = std::cmp::max(deepest_level, id.level());
        nodes_to_subsample.push(id);
    }
    let mut finished_nodes = HashMap::new();

    // sub sampling returns the list of finished nodes including all meta data
    // We start on the deepest level and work our way up the tree.
    for current_level in (0..deepest_level + 1).rev() {
        // All nodes on the same level can be subsampled in parallel.
        let res = nodes_to_subsample
            .into_iter()
            .partition(|n| n.level() == current_level);
        nodes_to_subsample = res.1;

        let mut parent_ids = HashSet::new();
        let mut subsample_nodes = Vec::new();
        for id in res.0 {
            let maybe_parent_id = id.parent_id();
            if maybe_parent_id.is_none() {
                // Only the root has no parents.
                assert_eq!(id.level(), 0);
                continue;
            }
            let parent_id = maybe_parent_id.unwrap();
            if parent_ids.contains(&parent_id) {
                continue;
            }
            parent_ids.insert(parent_id);
            subsample_nodes.push(parent_id);
        }

        let (finished_nodes_sender, finished_nodes_receiver) = mpsc::channel();
        pool.scoped(|scope| {
            for id in &subsample_nodes {
                let finished_nodes_sender_clone = finished_nodes_sender.clone();
                scope.execute(move || {
                    subsample_children_into(octree_meta, id, &finished_nodes_sender_clone).unwrap();
                });
            }
        });

        drop(finished_nodes_sender);
        for (id, num_points) in finished_nodes_receiver {
            if num_points > 0 {
                finished_nodes.insert(id, num_points);
            }
        }

        // The nodes that were just now created through sub-sampling will be required to create
        // their parents.
        nodes_to_subsample.extend(subsample_nodes.into_iter());
    }

    // Add all non-zero node meta data to meta.pb
    for (id, num_points) in finished_nodes {
        let bounding_cube = id.find_bounding_cube(&Cube::bounding(&octree_meta.bounding_box));
        let position_encoding =
            octree::PositionEncoding::new(&bounding_cube, octree_meta.resolution);

        let mut proto = proto::Node::new();
        proto.mut_id().set_level(id.level() as i32);
        proto.mut_id().set_index(id.index() as i64);
        proto.set_num_points(num_points);
        proto.set_position_encoding(position_encoding.to_proto());
        meta.mut_nodes().push(proto);
    }

    let mut buf_writer = BufWriter::new(File::create(&output_directory.join("meta.pb")).unwrap());
    meta.write_to_writer(&mut buf_writer).unwrap();
}
<|endoftext|>


<commit_before>//! Runs rustfmt on the repository.

use crate::Build;
use build_helper::t;
use ignore::WalkBuilder;
use std::path::Path;
use std::process::Command;

fn rustfmt(src: &Path, rustfmt: &Path, path: &Path, check: bool) {
    let mut cmd = Command::new(&rustfmt);
    // avoid the submodule config paths from coming into play,
    // we only allow a single global config for the workspace for now
    cmd.arg("--config-path").arg(&src.canonicalize().unwrap());
    cmd.arg("--edition").arg("2018");
    cmd.arg("--unstable-features");
    cmd.arg("--skip-children");
    if check {
        cmd.arg("--check");
    }
    cmd.arg(&path);
    let cmd_debug = format!("{:?}", cmd);
    let status = cmd.status().expect("executing rustfmt");
    if !status.success() {
        eprintln!(
            "Running `{}` failed.\nIf you're running `tidy`, \
            try again with `--bless` flag. Or, you just want to format \
            code, run `./x.py fmt` instead.",
            cmd_debug,
        );
        std::process::exit(1);
    }
}

#[derive(serde::Deserialize)]
struct RustfmtConfig {
    ignore: Vec<String>,
}

pub fn format(build: &Build, check: bool) {
    let mut builder = ignore::types::TypesBuilder::new();
    builder.add_defaults();
    builder.select("rust");
    let matcher = builder.build().unwrap();
    let rustfmt_config = build.src.join("rustfmt.toml");
    if !rustfmt_config.exists() {
        eprintln!("Not running formatting checks; rustfmt.toml does not exist.");
        eprintln!("This may happen in distributed tarballs.");
        return;
    }
    let rustfmt_config = t!(std::fs::read_to_string(&rustfmt_config));
    let rustfmt_config: RustfmtConfig = t!(toml::from_str(&rustfmt_config));
    let mut ignore_fmt = ignore::overrides::OverrideBuilder::new(&build.src);
    for ignore in rustfmt_config.ignore {
        ignore_fmt.add(&format!("!{}", ignore)).expect(&ignore);
    }
    let ignore_fmt = ignore_fmt.build().unwrap();

    let rustfmt_path = build.config.initial_rustfmt.as_ref().unwrap_or_else(|| {
        eprintln!("./x.py fmt is not supported on this channel");
        std::process::exit(1);
    });
    let src = build.src.clone();
    let walker = WalkBuilder::new(&build.src).types(matcher).overrides(ignore_fmt).build_parallel();
    walker.run(|| {
        let src = src.clone();
        let rustfmt_path = rustfmt_path.clone();
        Box::new(move |entry| {
            let entry = t!(entry);
            if entry.file_type().map_or(false, |t| t.is_file()) {
                rustfmt(&src, &rustfmt_path, &entry.path(), check);
            }
            ignore::WalkState::Continue
        })
    });
}
<commit_msg>Rollup merge of #69527 - pnkfelix:issue-69291-dont-run-rustfmt-on-untracked-paths, r=Mark-Simulacrum<commit_after>//! Runs rustfmt on the repository.

use crate::Build;
use build_helper::{output, t};
use ignore::WalkBuilder;
use std::path::Path;
use std::process::Command;

fn rustfmt(src: &Path, rustfmt: &Path, path: &Path, check: bool) {
    let mut cmd = Command::new(&rustfmt);
    // avoid the submodule config paths from coming into play,
    // we only allow a single global config for the workspace for now
    cmd.arg("--config-path").arg(&src.canonicalize().unwrap());
    cmd.arg("--edition").arg("2018");
    cmd.arg("--unstable-features");
    cmd.arg("--skip-children");
    if check {
        cmd.arg("--check");
    }
    cmd.arg(&path);
    let cmd_debug = format!("{:?}", cmd);
    let status = cmd.status().expect("executing rustfmt");
    if !status.success() {
        eprintln!(
            "Running `{}` failed.\nIf you're running `tidy`, \
            try again with `--bless` flag. Or, you just want to format \
            code, run `./x.py fmt` instead.",
            cmd_debug,
        );
        std::process::exit(1);
    }
}

#[derive(serde::Deserialize)]
struct RustfmtConfig {
    ignore: Vec<String>,
}

pub fn format(build: &Build, check: bool) {
    let mut builder = ignore::types::TypesBuilder::new();
    builder.add_defaults();
    builder.select("rust");
    let matcher = builder.build().unwrap();
    let rustfmt_config = build.src.join("rustfmt.toml");
    if !rustfmt_config.exists() {
        eprintln!("Not running formatting checks; rustfmt.toml does not exist.");
        eprintln!("This may happen in distributed tarballs.");
        return;
    }
    let rustfmt_config = t!(std::fs::read_to_string(&rustfmt_config));
    let rustfmt_config: RustfmtConfig = t!(toml::from_str(&rustfmt_config));
    let mut ignore_fmt = ignore::overrides::OverrideBuilder::new(&build.src);
    for ignore in rustfmt_config.ignore {
        ignore_fmt.add(&format!("!{}", ignore)).expect(&ignore);
    }
    let untracked_paths_output = output(
        Command::new("git").arg("status").arg("--porcelain").arg("--untracked-files=normal"),
    );
    let untracked_paths = untracked_paths_output
        .lines()
        .filter(|entry| entry.starts_with("??"))
        .map(|entry| entry.split(" ").nth(1).expect("every git status entry should list a path"));
    for untracked_path in untracked_paths {
        eprintln!("skip untracked path {} during rustfmt invocations", untracked_path);
        ignore_fmt.add(&format!("!{}", untracked_path)).expect(&untracked_path);
    }
    let ignore_fmt = ignore_fmt.build().unwrap();

    let rustfmt_path = build.config.initial_rustfmt.as_ref().unwrap_or_else(|| {
        eprintln!("./x.py fmt is not supported on this channel");
        std::process::exit(1);
    });
    let src = build.src.clone();
    let walker = WalkBuilder::new(&build.src).types(matcher).overrides(ignore_fmt).build_parallel();
    walker.run(|| {
        let src = src.clone();
        let rustfmt_path = rustfmt_path.clone();
        Box::new(move |entry| {
            let entry = t!(entry);
            if entry.file_type().map_or(false, |t| t.is_file()) {
                rustfmt(&src, &rustfmt_path, &entry.path(), check);
            }
            ignore::WalkState::Continue
        })
    });
}
<|endoftext|>