<commit_before>//! Handle network connections for a varlink service

#![allow(dead_code)]

#[cfg(unix)]
use libc::{close, dup2, getpid};
use std::io::{Read, Write};
use std::net::{Shutdown, TcpStream};
use std::process::Child;
use tempfile::TempDir;

#[cfg(unix)]
use std::os::unix::io::IntoRawFd;
#[cfg(unix)]
use std::os::unix::net::UnixStream;
#[cfg(windows)]
use uds_windows::UnixStream;

use crate::error::*;

pub enum VarlinkStream {
    TCP(TcpStream),
    UNIX(UnixStream),
}

#[cfg(windows)]
pub fn varlink_exec<S: ?Sized + AsRef<str>>(
    _address: &S,
) -> Result<(Child, String, Option<TempDir>)> {
    return Err(context!(ErrorKind::MethodNotImplemented(
        "varlink_exec".into()
    )));
}

#[cfg(unix)]
pub fn varlink_exec<S: ?Sized + AsRef<str>>(
    address: &S,
) -> Result<(Child, String, Option<TempDir>)> {
    use std::env;
    use std::os::unix::process::CommandExt;
    use std::process::Command;
    use tempfile::tempdir;

    let executable = String::from("exec ") + address.as_ref();

    use unix_socket::UnixListener;

    let dir = tempdir().map_err(map_context!())?;
    let file_path = dir.path().join("varlink-socket");

    let listener = UnixListener::bind(file_path.clone()).map_err(map_context!())?;
    let fd = listener.into_raw_fd();

    let child = Command::new("sh")
        .arg("-c")
        .arg(executable)
        .before_exec({
            let file_path = file_path.clone();
            move || {
                unsafe {
                    if fd != 3 {
                        close(3);
                        dup2(fd, 3);
                    }
                    env::set_var("VARLINK_ADDRESS", format!("unix:{}", file_path.display()));
                    env::set_var("LISTEN_FDS", "1");
                    env::set_var("LISTEN_FDNAMES", "varlink");
                    env::set_var("LISTEN_PID", format!("{}", getpid()));
                }
                Ok(())
            }
        })
        .spawn()
        .map_err(map_context!())?;
    Ok((child, format!("unix:{}", file_path.display()), Some(dir)))
}

#[cfg(windows)]
pub fn varlink_bridge<S: ?Sized + AsRef<str>>(address: &S) -> Result<(Child, VarlinkStream)> {
    use std::io::copy;
    use std::process::{Command, Stdio};
    use std::thread;

    let (stream0, stream1) = UnixStream::pair().map_err(map_context!())?;
    let executable = address.as_ref();

    let mut child = Command::new("cmd")
        .arg("/C")
        .arg(executable)
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .spawn()
        .map_err(map_context!())?;

    let mut client_writer = child.stdin.take().unwrap();
    let mut client_reader = child.stdout.take().unwrap();
    let mut service_writer = stream1.try_clone().map_err(map_context!())?;
    let mut service_reader = stream1;

    thread::spawn(move || copy(&mut client_reader, &mut service_writer));
    thread::spawn(move || copy(&mut service_reader, &mut client_writer));

    Ok((child, VarlinkStream::UNIX(stream0)))
}

#[cfg(unix)]
pub fn varlink_bridge<S: ?Sized + AsRef<str>>(address: &S) -> Result<(Child, VarlinkStream)> {
    use std::os::unix::io::FromRawFd;
    use std::process::Command;

    let executable = address.as_ref();
    // use unix_socket::UnixStream;
    let (stream0, stream1) = UnixStream::pair().map_err(map_context!())?;
    let fd = stream1.into_raw_fd();
    let childin = unsafe { ::std::fs::File::from_raw_fd(fd) };
    let childout = unsafe { ::std::fs::File::from_raw_fd(fd) };

    let child = Command::new("sh")
        .arg("-c")
        .arg(executable)
        .stdin(childin)
        .stdout(childout)
        .spawn()
        .map_err(map_context!())?;
    Ok((child, VarlinkStream::UNIX(stream0)))
}

#[cfg(any(target_os = "linux", target_os = "android"))]
fn get_abstract_unixstream(addr: &str) -> Result<UnixStream> {
    // FIXME: abstract unix domains sockets still not in std
    // FIXME: https://github.com/rust-lang/rust/issues/14194
    use std::os::unix::io::FromRawFd;
    use unix_socket::UnixStream as AbstractStream;

    unsafe {
        Ok(UnixStream::from_raw_fd(
            AbstractStream::connect(addr)
                .map_err(map_context!())?
                .into_raw_fd(),
        ))
    }
}

#[cfg(not(any(target_os = "linux", target_os = "android")))]
fn get_abstract_unixstream(_addr: &str) -> Result<UnixStream> {
    Err(context!(ErrorKind::InvalidAddress))
}

impl<'a> VarlinkStream {
    pub fn connect<S: ?Sized + AsRef<str>>(address: &S) -> Result<(Self, String)> {
        let address = address.as_ref();
        let new_address: String = address.into();

        if new_address.starts_with("tcp:") {
            Ok((
                VarlinkStream::TCP(
                    TcpStream::connect(&new_address[4..]).map_err(map_context!())?,
                ),
                new_address,
            ))
        } else if new_address.starts_with("unix:") {
            let mut addr = String::from(new_address[5..].split(';').next().unwrap());
            if addr.starts_with('@') {
                addr = addr.replacen('@', "\0", 1);
                return get_abstract_unixstream(&addr)
                    .and_then(|v| Ok((VarlinkStream::UNIX(v), new_address)));
            }
            Ok((
                VarlinkStream::UNIX(UnixStream::connect(addr).map_err(map_context!())?),
                new_address,
            ))
        } else {
            Err(context!(ErrorKind::InvalidAddress))?
        }
    }

    pub fn split(&mut self) -> Result<(Box<Read + Send + Sync>, Box<Write + Send + Sync>)> {
        match *self {
            VarlinkStream::TCP(ref mut s) => Ok((
                Box::new(s.try_clone().map_err(map_context!())?),
                Box::new(s.try_clone().map_err(map_context!())?),
            )),
            VarlinkStream::UNIX(ref mut s) => Ok((
                Box::new(s.try_clone().map_err(map_context!())?),
                Box::new(s.try_clone().map_err(map_context!())?),
            )),
        }
    }

    pub fn shutdown(&mut self) -> Result<()> {
        match *self {
            VarlinkStream::TCP(ref mut s) => s
                .shutdown(Shutdown::Both)
                .map_err(map_context!())?,
            VarlinkStream::UNIX(ref mut s) => s
                .shutdown(Shutdown::Both)
                .map_err(map_context!())?,
        }
        Ok(())
    }

    pub fn set_nonblocking(&self, b: bool) -> Result<()> {
        match *self {
            VarlinkStream::TCP(ref l) => l.set_nonblocking(b).map_err(map_context!())?,
            VarlinkStream::UNIX(ref l) => l.set_nonblocking(b).map_err(map_context!())?,
        }
        Ok(())
    }
}

impl Drop for VarlinkStream {
    fn drop(&mut self) {
        let _r = self.shutdown();
    }
}
<commit_msg>varlink/src/client.rs: implement AsRawFd<commit_after>//! Handle network connections for a varlink service

#![allow(dead_code)]

use std::io::{Read, Write};
use std::net::{Shutdown, TcpStream};
#[cfg(unix)]
use std::os::unix::io::IntoRawFd;
#[cfg(unix)]
use std::os::unix::io::{AsRawFd, RawFd};
#[cfg(unix)]
use std::os::unix::net::UnixStream;
use std::process::Child;

#[cfg(unix)]
use libc::{close, dup2, getpid};
use tempfile::TempDir;
#[cfg(windows)]
use uds_windows::UnixStream;

use crate::error::*;

pub enum VarlinkStream {
    TCP(TcpStream),
    UNIX(UnixStream),
}

#[cfg(windows)]
pub fn varlink_exec<S: ?Sized + AsRef<str>>(
    _address: &S,
) -> Result<(Child, String, Option<TempDir>)> {
    return Err(context!(ErrorKind::MethodNotImplemented(
        "varlink_exec".into()
    )));
}

#[cfg(unix)]
pub fn varlink_exec<S: ?Sized + AsRef<str>>(
    address: &S,
) -> Result<(Child, String, Option<TempDir>)> {
    use std::env;
    use std::os::unix::process::CommandExt;
    use std::process::Command;
    use tempfile::tempdir;

    let executable = String::from("exec ") + address.as_ref();

    use unix_socket::UnixListener;

    let dir = tempdir().map_err(map_context!())?;
    let file_path = dir.path().join("varlink-socket");

    let listener = UnixListener::bind(file_path.clone()).map_err(map_context!())?;
    let fd = listener.into_raw_fd();

    let child = Command::new("sh")
        .arg("-c")
        .arg(executable)
        .before_exec({
            let file_path = file_path.clone();
            move || {
                unsafe {
                    dup2(2, 1);
                    if fd != 3 {
                        close(3);
                        dup2(fd, 3);
                    }
                    env::set_var("VARLINK_ADDRESS", format!("unix:{}", file_path.display()));
                    env::set_var("LISTEN_FDS", "1");
                    env::set_var("LISTEN_FDNAMES", "varlink");
                    env::set_var("LISTEN_PID", format!("{}", getpid()));
                }
                Ok(())
            }
        })
        .spawn()
        .map_err(map_context!())?;
    Ok((child, format!("unix:{}", file_path.display()), Some(dir)))
}

#[cfg(windows)]
pub fn varlink_bridge<S: ?Sized + AsRef<str>>(address: &S) -> Result<(Child, VarlinkStream)> {
    use std::io::copy;
    use std::process::{Command, Stdio};
    use std::thread;

    let (stream0, stream1) = UnixStream::pair().map_err(map_context!())?;
    let executable = address.as_ref();

    let mut child = Command::new("cmd")
        .arg("/C")
        .arg(executable)
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .spawn()
        .map_err(map_context!())?;

    let mut client_writer = child.stdin.take().unwrap();
    let mut client_reader = child.stdout.take().unwrap();
    let mut service_writer = stream1.try_clone().map_err(map_context!())?;
    let mut service_reader = stream1;

    thread::spawn(move || copy(&mut client_reader, &mut service_writer));
    thread::spawn(move || copy(&mut service_reader, &mut client_writer));

    Ok((child, VarlinkStream::UNIX(stream0)))
}

#[cfg(unix)]
pub fn varlink_bridge<S: ?Sized + AsRef<str>>(address: &S) -> Result<(Child, VarlinkStream)> {
    use std::os::unix::io::FromRawFd;
    use std::process::Command;

    let executable = address.as_ref();
    // use unix_socket::UnixStream;
    let (stream0, stream1) = UnixStream::pair().map_err(map_context!())?;
    let fd = stream1.into_raw_fd();
    let childin = unsafe { ::std::fs::File::from_raw_fd(fd) };
    let childout = unsafe { ::std::fs::File::from_raw_fd(fd) };

    let child = Command::new("sh")
        .arg("-c")
        .arg(executable)
        .stdin(childin)
        .stdout(childout)
        .spawn()
        .map_err(map_context!())?;
    Ok((child, VarlinkStream::UNIX(stream0)))
}

#[cfg(any(target_os = "linux", target_os = "android"))]
fn get_abstract_unixstream(addr: &str) -> Result<UnixStream> {
    // FIXME: abstract unix domains sockets still not in std
    // FIXME: https://github.com/rust-lang/rust/issues/14194
    use std::os::unix::io::FromRawFd;
    use unix_socket::UnixStream as AbstractStream;

    unsafe {
        Ok(UnixStream::from_raw_fd(
            AbstractStream::connect(addr)
                .map_err(map_context!())?
                .into_raw_fd(),
        ))
    }
}

#[cfg(not(any(target_os = "linux", target_os = "android")))]
fn get_abstract_unixstream(_addr: &str) -> Result<UnixStream> {
    Err(context!(ErrorKind::InvalidAddress))
}

impl<'a> VarlinkStream {
    pub fn connect<S: ?Sized + AsRef<str>>(address: &S) -> Result<(Self, String)> {
        let address = address.as_ref();
        let new_address: String = address.into();

        if new_address.starts_with("tcp:") {
            Ok((
                VarlinkStream::TCP(
                    TcpStream::connect(&new_address[4..]).map_err(map_context!())?,
                ),
                new_address,
            ))
        } else if new_address.starts_with("unix:") {
            let mut addr = String::from(new_address[5..].split(';').next().unwrap());
            if addr.starts_with('@') {
                addr = addr.replacen('@', "\0", 1);
                return get_abstract_unixstream(&addr)
                    .and_then(|v| Ok((VarlinkStream::UNIX(v), new_address)));
            }
            Ok((
                VarlinkStream::UNIX(UnixStream::connect(addr).map_err(map_context!())?),
                new_address,
            ))
        } else {
            Err(context!(ErrorKind::InvalidAddress))?
        }
    }

    pub fn split(&mut self) -> Result<(Box<Read + Send + Sync>, Box<Write + Send + Sync>)> {
        match *self {
            VarlinkStream::TCP(ref mut s) => Ok((
                Box::new(s.try_clone().map_err(map_context!())?),
                Box::new(s.try_clone().map_err(map_context!())?),
            )),
            VarlinkStream::UNIX(ref mut s) => Ok((
                Box::new(s.try_clone().map_err(map_context!())?),
                Box::new(s.try_clone().map_err(map_context!())?),
            )),
        }
    }

    pub fn shutdown(&mut self) -> Result<()> {
        match *self {
            VarlinkStream::TCP(ref mut s) => s
                .shutdown(Shutdown::Both)
                .map_err(map_context!())?,
            VarlinkStream::UNIX(ref mut s) => s
                .shutdown(Shutdown::Both)
                .map_err(map_context!())?,
        }
        Ok(())
    }

    pub fn set_nonblocking(&self, b: bool) -> Result<()> {
        match *self {
            VarlinkStream::TCP(ref l) => l.set_nonblocking(b).map_err(map_context!())?,
            VarlinkStream::UNIX(ref l) => l.set_nonblocking(b).map_err(map_context!())?,
        }
        Ok(())
    }
}

impl AsRawFd for VarlinkStream {
    fn as_raw_fd(&self) -> RawFd {
        match *self {
            VarlinkStream::TCP(ref l) => l.as_raw_fd(),
            VarlinkStream::UNIX(ref l) => l.as_raw_fd(),
        }
    }
}

impl Drop for VarlinkStream {
    fn drop(&mut self) {
        let _r = self.shutdown();
    }
}
<|endoftext|>


<commit_before>use std::sync::mpsc::Sender;

use anyhow::{bail, Result};
use async_std::task;
use log::{debug, info, warn};

use pueue::message::*;
use pueue::protocol::*;
use pueue::state::SharedState;

use crate::cli::CliArguments;
use crate::instructions::handle_message;
use crate::streaming::handle_follow;

/// Poll the unix listener and accept new incoming connections.
/// Create a new future to handle the message and spawn it.
pub async fn accept_incoming(sender: Sender<Message>, state: SharedState, opt: CliArguments) -> Result<()> {
    let (unix_socket_path, port) = {
        let state = state.lock().unwrap();
        let shared = &state.settings.shared;

        // Return the unix socket path, if we're supposed to use it.
        if shared.use_unix_socket {
            (Some(shared.unix_socket_path.clone()), None)
        } else {
            // Otherwise use tcp sockets and a given port
            // Commandline argument overwrites the configuration files values for port.
            let port = if let Some(port) = opt.port.clone() {
                port
            } else {
                shared.port.clone()
            };

            (None, Some(port))
        }
    };

    let listener = get_listener(unix_socket_path, port).await?;

    loop {
        // Poll if we have a new incoming connection.
        let socket = listener.accept().await?;

        // Start a new task for the request
        let sender_clone = sender.clone();
        let state_clone = state.clone();
        task::spawn(async move {
            let _result = handle_incoming(socket, sender_clone, state_clone).await;
        });
    }
}

/// Continuously poll the existing incoming futures.
/// In case we received an instruction, handle it and create a response future.
/// The response future is added to unix_responses and handled in a separate function.
async fn handle_incoming(
    mut socket: Socket,
    sender: Sender<Message>,
    state: SharedState,
) -> Result<()> {
    // Receive the secret once and check, whether the client is allowed to connect
    let payload_bytes = receive_bytes(&mut socket).await?;

    // Didn't receive any bytes. The client disconnected.
    if payload_bytes.is_empty() {
        info!("Client went away");
        return Ok(());
    }

    let secret = String::from_utf8(payload_bytes)?;

    // Return immediately, if we got a wrong secret from the client.
    {
        let state = state.lock().unwrap();
        if secret != state.settings.shared.secret {
            warn!("Received invalid secret: {}", secret);
            bail!("Received invalid secret");
        }
    }
    // Send a super short `ok` byte to the client, so it knows that the secret has been accepted.
    send_bytes(b"hello", &mut socket).await?;

    // Save the directory for convenience purposes and to prevent continuously
    // locking the state in the streaming loop.
    let pueue_directory = {
        let state = state.lock().unwrap();
        state.settings.shared.pueue_directory.clone()
    };

    loop {
        // Receive the actual instruction from the client
        let message = receive_message(&mut socket).await?;
        debug!("Received instruction: {:?}", message);

        let response = if let Message::StreamRequest(message) = message {
            // The client requested the output of a task.
            // Since we allow streaming, this needs to be handled seperately.
            handle_follow(&pueue_directory, &mut socket, &state, message).await?
        } else {
            // Process a normal message.
            handle_message(message, &sender, &state)
        };

        // Respond to the client.
        send_message(response, &mut socket).await?;
    }
}
<commit_msg>Cargo fmt<commit_after>use std::sync::mpsc::Sender;

use anyhow::{bail, Result};
use async_std::task;
use log::{debug, info, warn};

use pueue::message::*;
use pueue::protocol::*;
use pueue::state::SharedState;

use crate::cli::CliArguments;
use crate::instructions::handle_message;
use crate::streaming::handle_follow;

/// Poll the unix listener and accept new incoming connections.
/// Create a new future to handle the message and spawn it.
pub async fn accept_incoming(
    sender: Sender<Message>,
    state: SharedState,
    opt: CliArguments,
) -> Result<()> {
    let (unix_socket_path, port) = {
        let state = state.lock().unwrap();
        let shared = &state.settings.shared;

        // Return the unix socket path, if we're supposed to use it.
        if shared.use_unix_socket {
            (Some(shared.unix_socket_path.clone()), None)
        } else {
            // Otherwise use tcp sockets and a given port
            // Commandline argument overwrites the configuration files values for port.
            let port = if let Some(port) = opt.port.clone() {
                port
            } else {
                shared.port.clone()
            };

            (None, Some(port))
        }
    };

    let listener = get_listener(unix_socket_path, port).await?;

    loop {
        // Poll if we have a new incoming connection.
        let socket = listener.accept().await?;

        // Start a new task for the request
        let sender_clone = sender.clone();
        let state_clone = state.clone();
        task::spawn(async move {
            let _result = handle_incoming(socket, sender_clone, state_clone).await;
        });
    }
}

/// Continuously poll the existing incoming futures.
/// In case we received an instruction, handle it and create a response future.
/// The response future is added to unix_responses and handled in a separate function.
async fn handle_incoming(
    mut socket: Socket,
    sender: Sender<Message>,
    state: SharedState,
) -> Result<()> {
    // Receive the secret once and check, whether the client is allowed to connect
    let payload_bytes = receive_bytes(&mut socket).await?;

    // Didn't receive any bytes. The client disconnected.
    if payload_bytes.is_empty() {
        info!("Client went away");
        return Ok(());
    }

    let secret = String::from_utf8(payload_bytes)?;

    // Return immediately, if we got a wrong secret from the client.
    {
        let state = state.lock().unwrap();
        if secret != state.settings.shared.secret {
            warn!("Received invalid secret: {}", secret);
            bail!("Received invalid secret");
        }
    }
    // Send a super short `ok` byte to the client, so it knows that the secret has been accepted.
    send_bytes(b"hello", &mut socket).await?;

    // Save the directory for convenience purposes and to prevent continuously
    // locking the state in the streaming loop.
    let pueue_directory = {
        let state = state.lock().unwrap();
        state.settings.shared.pueue_directory.clone()
    };

    loop {
        // Receive the actual instruction from the client
        let message = receive_message(&mut socket).await?;
        debug!("Received instruction: {:?}", message);

        let response = if let Message::StreamRequest(message) = message {
            // The client requested the output of a task.
            // Since we allow streaming, this needs to be handled seperately.
            handle_follow(&pueue_directory, &mut socket, &state, message).await?
        } else {
            // Process a normal message.
            handle_message(message, &sender, &state)
        };

        // Respond to the client.
        send_message(response, &mut socket).await?;
    }
}
<|endoftext|>


<commit_before>//! Contains implementation of Binary FBX parser.

extern crate byteorder;
extern crate flate2;

use std::io::Read;
use reader::error::{Result, Error, ErrorKind};
use reader::FbxEvent;
use common::OwnedProperty;
use super::CommonState;

/// A parser for Binary FBX.
#[derive(Debug, Clone)]
pub struct BinaryParser {
    version: u32,
    end_offset_stack: Vec<u32>,
}

impl BinaryParser {
    /// Constructs Binary FBX parser with FBX version (which is placed after magic binary).
    pub fn new(version: u32) -> Self {
        BinaryParser {
            version: version,
            end_offset_stack: vec![],
        }
    }

    pub fn next<R: Read>(&mut self, reader: &mut R, common: &mut CommonState) -> Result<FbxEvent> {
        // Check if the previously read node ends here.
        if let Some(&end_pos_top) = self.end_offset_stack.last() {
            if end_pos_top as u64 == common.pos {
                // Reached the end of previously read node.
                self.end_offset_stack.pop();
                return Ok(FbxEvent::EndNode);
            }
        }

        // Read a node record header.
        let node_record_header = try!(NodeRecordHeader::read(reader, &mut common.pos));
        if node_record_header.is_null_record() {
            // End of a node.
            return if let Some(expected_pos) = self.end_offset_stack.pop() {
                if common.pos == expected_pos as u64 {
                    Ok(FbxEvent::EndNode)
                } else {
                    // Data is collapsed (the node doesn't end at expected position).
                    Err(Error::new(
                            common.pos,
                            ErrorKind::DataError(format!("Node does not end at expected position (expected {}, now at {})", expected_pos, common.pos))))
                }
            } else {
                // Reached end of all nodes.
                // (Extra NULL-record header is end marker of implicit root node.)
                // Footer with unknown contents follows.
                // TODO: Read footer.
                //       Files exported by official products or SDK have padding and their file
                //       sizes are multiple of 16, but some files exported by third-party apps
                //       (such as blender) does not.
                //       So it may be difficult to check if the footer is correct or wrong.
                // NOTE: There is the only thing known, the last 16 bytes of the data always seem
                //       to be `[0xf8, 0x5a, 0x8c, 0x6a, 0xde, 0xf5, 0xd9, 0x7e, 0xec, 0xe9, 0x0c,
                //       0xe3, 0x75, 0x8f, 0x29, 0x0b]`.
                Ok(FbxEvent::EndFbx)
            };
        } else {
            // Start of a node.
            self.end_offset_stack.push(node_record_header.end_offset);
        }

        // Read a node name.
        let name = try_read_fixstr!(common.pos, reader, node_record_header.name_len);

        // Read properties.
        let mut properties = Vec::<OwnedProperty>::with_capacity(node_record_header.num_properties as usize);
        for _ in 0..node_record_header.num_properties {
            let prop = try!(self.read_property(reader, common));
            properties.push(prop);
        }

        Ok(FbxEvent::StartNode {
            name: name,
            properties: properties,
        })
    }

    /// Read a node property value.
    fn read_property<R: Read>(&mut self, reader: &mut R, common: &mut CommonState) -> Result<OwnedProperty> {
        let type_code = try_read_le_u8!(common.pos, reader);
        // type code must be ASCII.
        let type_code = if type_code > 0x80 {
            return Err(Error::new(common.pos-1, ErrorKind::DataError(format!("Expected property type code (ASCII) but got {:#x}", type_code))));
        } else {
            type_code as char
        };
        let value = match type_code {
            // 1 bit boolean (1: true, 0: false) encoded as the LSB of a 1 byte value.
            'C' => {
                let val = try_read_le_u8!(common.pos, reader);
                // It seems 'T' (0x54) is used as `false`, 'T' (0x59) is used as `true`.
                if (val != 'T' as u8) && (val != 'Y' as u8) {
                    // Should this treated as error?
                    // (I don't know whether other characters than 'T' and 'Y' are allowed...)
                    warn!("Expected 'T' or 'Y' for representaton of boolean property value, but got {:#x}", val);
                }
                // Check LSB.
                OwnedProperty::Bool(val & 1 == 1)
            },
            // 2 byte signed integer.
            'Y' => {
                OwnedProperty::I16(try_read_le_i16!(common.pos, reader))
            },
            // 4 byte signed integer.
            'I' => {
                OwnedProperty::I32(try_read_le_i32!(common.pos, reader))
            },
            // 4 byte single-precision IEEE 754 floating-point number.
            'F' => {
                OwnedProperty::F32(try_read_le_f32!(common.pos, reader))
            },
            // 8 byte double-precision IEEE 754 floating-point number.
            'D' => {
                OwnedProperty::F64(try_read_le_f64!(common.pos, reader))
            },
            // 8 byte signed integer.
            'L' => {
                OwnedProperty::I64(try_read_le_i64!(common.pos, reader))
            },
            // Array types
            'f'|'d'|'l'|'i'|'b' => {
                let array_header = try!(PropertyArrayHeader::read(reader, &mut common.pos));
                try!(self.read_property_value_array(reader, common, type_code, &array_header))
            },
            // String
            'S' => {
                let length = try_read_le_u32!(common.pos, reader);
                OwnedProperty::String(try_read_fixstr!(common.pos, reader, length))
            },
            // Raw binary data
            'R' => {
                let length = try_read_le_u32!(common.pos, reader);
                OwnedProperty::Binary(try_read_exact!(common.pos, reader, length))
            },
            _ => {
                return Err(Error::new(
                        common.pos,
                        ErrorKind::UnexpectedValue(format!(
                                "Unsupported type code appears in node property: type_code={}({:#x})",
                                type_code, type_code as u8))));
            }
        };
        Ok(value)
    }

    /// Read a property value of array type from given stream which maybe compressed.
    fn read_property_value_array<R: Read>(&mut self,
                                          reader: &mut R, common: &mut CommonState,
                                          type_code: char, array_header: &PropertyArrayHeader) -> Result<OwnedProperty> {
        match array_header.encoding {
            // 0; raw
            0 => {
                let (val, byte_size) = try!(self.read_property_value_array_from_plain_stream(reader, common.pos, type_code, array_header.array_length));
                common.pos += byte_size;
                Ok(val)
            },
            // 1: zlib compressed data
            1 => {
                let mut decoded_stream = flate2::read::ZlibDecoder::new(reader.by_ref().take(array_header.compressed_length as u64));
                let (val, _) = try!(self.read_property_value_array_from_plain_stream(&mut decoded_stream, common.pos, type_code, array_header.array_length));
                common.pos += array_header.compressed_length as u64;
                Ok(val)
            },
            // Unknown.
            e => {
                Err(Error::new(
                        common.pos,
                        ErrorKind::UnexpectedValue(format!("Unsupported property array encoding, got {:#x}", e))))
            }
        }
    }

    /// Read a property value of array type from plain (uncompressed) stream.
    fn read_property_value_array_from_plain_stream<R: Read>(&mut self, reader: &mut R, abs_pos: u64, type_code: char,
                                                            num_elements: u32) -> Result<(OwnedProperty, u64)> {
        use self::byteorder::{ReadBytesExt, LittleEndian};
        Ok(match type_code {
            // Array of 4 byte single-precision IEEE 754 floating-point number.
            'f' => {
                let mut data = Vec::<f32>::with_capacity(num_elements as usize);
                for _ in 0..num_elements {
                    data.push(try_with_pos!(abs_pos, reader.read_f32::<LittleEndian>()));
                }
                (OwnedProperty::VecF32(data), num_elements as u64 * 4)
            },
            // Array of 8 byte double-precision IEEE 754 floating-point number.
            'd' => {
                let mut data = Vec::<f64>::with_capacity(num_elements as usize);
                for _ in 0..num_elements {
                    data.push(try_with_pos!(abs_pos, reader.read_f64::<LittleEndian>()));
                }
                (OwnedProperty::VecF64(data), num_elements as u64 * 8)
            },
            // Array of 8 byte signed integer.
            'l' => {
                let mut data = Vec::<i64>::with_capacity(num_elements as usize);
                for _ in 0..num_elements {
                    data.push(try_with_pos!(abs_pos, reader.read_i64::<LittleEndian>()));
                }
                (OwnedProperty::VecI64(data), num_elements as u64 * 8)
            },
            // Array of 4 byte signed integer.
            'i' => {
                let mut data = Vec::<i32>::with_capacity(num_elements as usize);
                for _ in 0..num_elements {
                    data.push(try_with_pos!(abs_pos, reader.read_i32::<LittleEndian>()));
                }
                (OwnedProperty::VecI32(data), num_elements as u64 * 4)
            },
            // Array of 1 byte booleans (always 0 or 1?).
            'b' => {
                let mut data = Vec::<bool>::with_capacity(num_elements as usize);
                for _ in 0..num_elements {
                    // Check LSB.
                    data.push(try_with_pos!(abs_pos, reader.read_u8()) & 1 == 1);
                }
                (OwnedProperty::VecBool(data), num_elements as u64)
            },
            _ => {
                // Unreachable because `read_property()` gives only 'f' , 'd', 'l', 'i', or 'b' to
                // `read_property_value_array()`.
                unreachable!();
            }
        })
    }
}

/// A header of a node.
#[derive(Debug, Copy, Clone)]
struct NodeRecordHeader {
    /// Position of the end of the node.
    end_offset: u32,
    /// Number of the properties the node has.
    num_properties: u32,
    /// Byte size of properties of the node in the FBX stream.
    property_list_len: u32,
    /// Byte size of the node name.
    name_len: u8,
}

impl NodeRecordHeader {
    /// Constructs `NodeRecordHeader` from the given stream.
    pub fn read<R: Read>(reader: &mut R, pos: &mut u64) -> Result<Self> {
        let end_offset = try_read_le_u32!(*pos, reader);
        let num_properties = try_read_le_u32!(*pos, reader);
        let property_list_len = try_read_le_u32!(*pos, reader);
        let name_len = try_read_le_u8!(*pos, reader);
        Ok(NodeRecordHeader {
            end_offset: end_offset,
            num_properties: num_properties,
            property_list_len: property_list_len,
            name_len: name_len,
        })
    }

    /// Check whether the header indicates there are no more children.
    pub fn is_null_record(&self) -> bool {
        self.end_offset == 0
            && self.num_properties == 0
            && self.property_list_len == 0
            && self.name_len == 0
    }
}

/// A header of a property of array type.
#[derive(Debug, Copy, Clone)]
pub struct PropertyArrayHeader {
    /// Number of values in the array, *NOT byte size*.
    array_length: u32,
    /// Denotes whether data in stream is plain, or what algorithm it is compressed by.
    encoding: u32,
    /// Byte size of the compressed array value in the stream.
    compressed_length: u32,
}

impl PropertyArrayHeader {
    /// Constructs `PropertyArrayHeader` from the given stream.
    pub fn read<R: Read>(reader: &mut R, pos: &mut u64) -> Result<Self> {
        let array_length = try_read_le_u32!(*pos, reader);
        let encoding = try_read_le_u32!(*pos, reader);
        let compressed_length = try_read_le_u32!(*pos, reader);
        Ok(PropertyArrayHeader {
            array_length: array_length,
            encoding: encoding,
            compressed_length: compressed_length,
        })
    }
}
<commit_msg>Support FBX 7.5 node record header<commit_after>//! Contains implementation of Binary FBX parser.

extern crate byteorder;
extern crate flate2;

use std::io::Read;
use reader::error::{Result, Error, ErrorKind};
use reader::FbxEvent;
use common::OwnedProperty;
use super::CommonState;

/// A parser for Binary FBX.
#[derive(Debug, Clone)]
pub struct BinaryParser {
    version: u32,
    end_offset_stack: Vec<u64>,
}

impl BinaryParser {
    /// Constructs Binary FBX parser with FBX version (which is placed after magic binary).
    pub fn new(version: u32) -> Self {
        BinaryParser {
            version: version,
            end_offset_stack: vec![],
        }
    }

    pub fn next<R: Read>(&mut self, reader: &mut R, common: &mut CommonState) -> Result<FbxEvent> {
        // Check if the previously read node ends here.
        if let Some(&end_pos_top) = self.end_offset_stack.last() {
            if end_pos_top as u64 == common.pos {
                // Reached the end of previously read node.
                self.end_offset_stack.pop();
                return Ok(FbxEvent::EndNode);
            }
        }

        // Read a node record header.
        let node_record_header = try!(NodeRecordHeader::read(reader, &mut common.pos, self));
        if node_record_header.is_null_record() {
            // End of a node.
            return if let Some(expected_pos) = self.end_offset_stack.pop() {
                if common.pos == expected_pos as u64 {
                    Ok(FbxEvent::EndNode)
                } else {
                    // Data is collapsed (the node doesn't end at expected position).
                    Err(Error::new(
                            common.pos,
                            ErrorKind::DataError(format!("Node does not end at expected position (expected {}, now at {})", expected_pos, common.pos))))
                }
            } else {
                // Reached end of all nodes.
                // (Extra NULL-record header is end marker of implicit root node.)
                // Footer with unknown contents follows.
                // TODO: Read footer.
                //       Files exported by official products or SDK have padding and their file
                //       sizes are multiple of 16, but some files exported by third-party apps
                //       (such as blender) does not.
                //       So it may be difficult to check if the footer is correct or wrong.
                // NOTE: There is the only thing known, the last 16 bytes of the data always seem
                //       to be `[0xf8, 0x5a, 0x8c, 0x6a, 0xde, 0xf5, 0xd9, 0x7e, 0xec, 0xe9, 0x0c,
                //       0xe3, 0x75, 0x8f, 0x29, 0x0b]`.
                Ok(FbxEvent::EndFbx)
            };
        } else {
            // Start of a node.
            self.end_offset_stack.push(node_record_header.end_offset);
        }

        // Read a node name.
        let name = try_read_fixstr!(common.pos, reader, node_record_header.name_len);

        // Read properties.
        let mut properties = Vec::<OwnedProperty>::with_capacity(node_record_header.num_properties as usize);
        for _ in 0..node_record_header.num_properties {
            let prop = try!(self.read_property(reader, common));
            properties.push(prop);
        }

        Ok(FbxEvent::StartNode {
            name: name,
            properties: properties,
        })
    }

    /// Read a node property value.
    fn read_property<R: Read>(&mut self, reader: &mut R, common: &mut CommonState) -> Result<OwnedProperty> {
        let type_code = try_read_le_u8!(common.pos, reader);
        // type code must be ASCII.
        let type_code = if type_code > 0x80 {
            return Err(Error::new(common.pos-1, ErrorKind::DataError(format!("Expected property type code (ASCII) but got {:#x}", type_code))));
        } else {
            type_code as char
        };
        let value = match type_code {
            // 1 bit boolean (1: true, 0: false) encoded as the LSB of a 1 byte value.
            'C' => {
                let val = try_read_le_u8!(common.pos, reader);
                // It seems 'T' (0x54) is used as `false`, 'T' (0x59) is used as `true`.
                if (val != 'T' as u8) && (val != 'Y' as u8) {
                    // Should this treated as error?
                    // (I don't know whether other characters than 'T' and 'Y' are allowed...)
                    warn!("Expected 'T' or 'Y' for representaton of boolean property value, but got {:#x}", val);
                }
                // Check LSB.
                OwnedProperty::Bool(val & 1 == 1)
            },
            // 2 byte signed integer.
            'Y' => {
                OwnedProperty::I16(try_read_le_i16!(common.pos, reader))
            },
            // 4 byte signed integer.
            'I' => {
                OwnedProperty::I32(try_read_le_i32!(common.pos, reader))
            },
            // 4 byte single-precision IEEE 754 floating-point number.
            'F' => {
                OwnedProperty::F32(try_read_le_f32!(common.pos, reader))
            },
            // 8 byte double-precision IEEE 754 floating-point number.
            'D' => {
                OwnedProperty::F64(try_read_le_f64!(common.pos, reader))
            },
            // 8 byte signed integer.
            'L' => {
                OwnedProperty::I64(try_read_le_i64!(common.pos, reader))
            },
            // Array types
            'f'|'d'|'l'|'i'|'b' => {
                let array_header = try!(PropertyArrayHeader::read(reader, &mut common.pos));
                try!(self.read_property_value_array(reader, common, type_code, &array_header))
            },
            // String
            'S' => {
                let length = try_read_le_u32!(common.pos, reader);
                OwnedProperty::String(try_read_fixstr!(common.pos, reader, length))
            },
            // Raw binary data
            'R' => {
                let length = try_read_le_u32!(common.pos, reader);
                OwnedProperty::Binary(try_read_exact!(common.pos, reader, length))
            },
            _ => {
                return Err(Error::new(
                        common.pos,
                        ErrorKind::UnexpectedValue(format!(
                                "Unsupported type code appears in node property: type_code={}({:#x})",
                                type_code, type_code as u8))));
            }
        };
        Ok(value)
    }

    /// Read a property value of array type from given stream which maybe compressed.
    fn read_property_value_array<R: Read>(&mut self,
                                          reader: &mut R, common: &mut CommonState,
                                          type_code: char, array_header: &PropertyArrayHeader) -> Result<OwnedProperty> {
        match array_header.encoding {
            // 0; raw
            0 => {
                let (val, byte_size) = try!(self.read_property_value_array_from_plain_stream(reader, common.pos, type_code, array_header.array_length));
                common.pos += byte_size;
                Ok(val)
            },
            // 1: zlib compressed data
            1 => {
                let mut decoded_stream = flate2::read::ZlibDecoder::new(reader.by_ref().take(array_header.compressed_length as u64));
                let (val, _) = try!(self.read_property_value_array_from_plain_stream(&mut decoded_stream, common.pos, type_code, array_header.array_length));
                common.pos += array_header.compressed_length as u64;
                Ok(val)
            },
            // Unknown.
            e => {
                Err(Error::new(
                        common.pos,
                        ErrorKind::UnexpectedValue(format!("Unsupported property array encoding, got {:#x}", e))))
            }
        }
    }

    /// Read a property value of array type from plain (uncompressed) stream.
    fn read_property_value_array_from_plain_stream<R: Read>(&mut self, reader: &mut R, abs_pos: u64, type_code: char,
                                                            num_elements: u32) -> Result<(OwnedProperty, u64)> {
        use self::byteorder::{ReadBytesExt, LittleEndian};
        Ok(match type_code {
            // Array of 4 byte single-precision IEEE 754 floating-point number.
            'f' => {
                let mut data = Vec::<f32>::with_capacity(num_elements as usize);
                for _ in 0..num_elements {
                    data.push(try_with_pos!(abs_pos, reader.read_f32::<LittleEndian>()));
                }
                (OwnedProperty::VecF32(data), num_elements as u64 * 4)
            },
            // Array of 8 byte double-precision IEEE 754 floating-point number.
            'd' => {
                let mut data = Vec::<f64>::with_capacity(num_elements as usize);
                for _ in 0..num_elements {
                    data.push(try_with_pos!(abs_pos, reader.read_f64::<LittleEndian>()));
                }
                (OwnedProperty::VecF64(data), num_elements as u64 * 8)
            },
            // Array of 8 byte signed integer.
            'l' => {
                let mut data = Vec::<i64>::with_capacity(num_elements as usize);
                for _ in 0..num_elements {
                    data.push(try_with_pos!(abs_pos, reader.read_i64::<LittleEndian>()));
                }
                (OwnedProperty::VecI64(data), num_elements as u64 * 8)
            },
            // Array of 4 byte signed integer.
            'i' => {
                let mut data = Vec::<i32>::with_capacity(num_elements as usize);
                for _ in 0..num_elements {
                    data.push(try_with_pos!(abs_pos, reader.read_i32::<LittleEndian>()));
                }
                (OwnedProperty::VecI32(data), num_elements as u64 * 4)
            },
            // Array of 1 byte booleans (always 0 or 1?).
            'b' => {
                let mut data = Vec::<bool>::with_capacity(num_elements as usize);
                for _ in 0..num_elements {
                    // Check LSB.
                    data.push(try_with_pos!(abs_pos, reader.read_u8()) & 1 == 1);
                }
                (OwnedProperty::VecBool(data), num_elements as u64)
            },
            _ => {
                // Unreachable because `read_property()` gives only 'f' , 'd', 'l', 'i', or 'b' to
                // `read_property_value_array()`.
                unreachable!();
            }
        })
    }
}

/// A header of a node.
#[derive(Debug, Copy, Clone)]
struct NodeRecordHeader {
    /// Position of the end of the node.
    end_offset: u64,
    /// Number of the properties the node has.
    num_properties: u64,
    /// Byte size of properties of the node in the FBX stream.
    property_list_len: u64,
    /// Byte size of the node name.
    name_len: u8,
}

impl NodeRecordHeader {
    /// Constructs `NodeRecordHeader` from the given stream.
    pub fn read<R: Read>(reader: &mut R, pos: &mut u64, context: &BinaryParser) -> Result<Self> {
        let end_offset = if context.version < 7500 {
            try_read_le_u32!(*pos, reader) as u64
        } else {
            try_read_le_u64!(*pos, reader)
        };
        let num_properties = if context.version < 7500 {
            try_read_le_u32!(*pos, reader) as u64
        } else {
            try_read_le_u64!(*pos, reader)
        };
        let property_list_len = if context.version < 7500 {
            try_read_le_u32!(*pos, reader) as u64
        } else {
            try_read_le_u64!(*pos, reader)
        };
        let name_len = try_read_le_u8!(*pos, reader);
        Ok(NodeRecordHeader {
            end_offset: end_offset,
            num_properties: num_properties,
            property_list_len: property_list_len,
            name_len: name_len,
        })
    }

    /// Check whether the header indicates there are no more children.
    pub fn is_null_record(&self) -> bool {
        self.end_offset == 0
            && self.num_properties == 0
            && self.property_list_len == 0
            && self.name_len == 0
    }
}

/// A header of a property of array type.
#[derive(Debug, Copy, Clone)]
pub struct PropertyArrayHeader {
    /// Number of values in the array, *NOT byte size*.
    array_length: u32,
    /// Denotes whether data in stream is plain, or what algorithm it is compressed by.
    encoding: u32,
    /// Byte size of the compressed array value in the stream.
    compressed_length: u32,
}

impl PropertyArrayHeader {
    /// Constructs `PropertyArrayHeader` from the given stream.
    pub fn read<R: Read>(reader: &mut R, pos: &mut u64) -> Result<Self> {
        let array_length = try_read_le_u32!(*pos, reader);
        let encoding = try_read_le_u32!(*pos, reader);
        let compressed_length = try_read_le_u32!(*pos, reader);
        Ok(PropertyArrayHeader {
            array_length: array_length,
            encoding: encoding,
            compressed_length: compressed_length,
        })
    }
}
<|endoftext|>


<commit_before>//! Reduce the collection to one occurrence of each distinct element.
//!
//! The `distinct_total` and `distinct_total_u` operators are optimizations of the more general
//! `distinct` and `distinct_u` operators for the case in which time is totally ordered.

use timely::order::TotalOrder;
use timely::dataflow::*;
use timely::dataflow::operators::Unary;
use timely::dataflow::channels::pact::Pipeline;

use lattice::Lattice;
use ::{Data, Collection, Diff};
use hashable::Hashable;
use collection::AsCollection;
use operators::arrange::{Arranged, ArrangeBySelf};
use trace::{BatchReader, Cursor, TraceReader};

/// Extension trait for the `distinct` differential dataflow method.
pub trait ThresholdTotal<G: Scope, K: Data, R: Diff> where G::Timestamp: TotalOrder+Lattice+Ord {
    /// Reduces the collection to one occurrence of each distinct element.
    ///
    /// # Examples
    ///
    /// ```
    /// extern crate timely;
    /// extern crate differential_dataflow;
    ///
    /// use differential_dataflow::input::Input;
    /// use differential_dataflow::operators::ThresholdTotal;
    ///
    /// fn main() {
    ///     ::timely::example(|scope| {
    ///         // report the number of occurrences of each key
    ///         scope.new_collection_from(1 .. 10).1
    ///              .map(|x| x / 3)
    ///              .threshold_total(|c| c % 2);
    ///     });
    /// }
    /// ```
    fn threshold_total<R2: Diff, F: Fn(R)->R2+'static>(&self, thresh: F) -> Collection<G, K, R2>;
    /// Reduces the collection to one occurrence of each distinct element.
    ///
    /// This reduction only tests whether the weight associated with a record is non-zero, and otherwise
    /// ignores its specific value. To take more general actions based on the accumulated weight, consider
    /// the `threshold` method.
    ///
    /// # Examples
    ///
    /// ```
    /// extern crate timely;
    /// extern crate differential_dataflow;
    ///
    /// use differential_dataflow::input::Input;
    /// use differential_dataflow::operators::ThresholdTotal;
    ///
    /// fn main() {
    ///     ::timely::example(|scope| {
    ///         // report the number of occurrences of each key
    ///         scope.new_collection_from(1 .. 10).1
    ///              .map(|x| x / 3)
    ///              .distinct_total();
    ///     });
    /// }
    /// ```
    fn distinct_total(&self) -> Collection<G, K, isize> {
        self.threshold_total(|c| if c.is_zero() { 0 } else { 1 })
    }
}

impl<G: Scope, K: Data+Hashable, R: Diff> ThresholdTotal<G, K, R> for Collection<G, K, R>
where G::Timestamp: TotalOrder+Lattice+Ord {
    fn threshold_total<R2: Diff, F: Fn(R)->R2+'static>(&self, thresh: F) -> Collection<G, K, R2> {
        self.arrange_by_self()
            .threshold_total(thresh)
    }
}

impl<G: Scope, K: Data, R: Diff, T1> ThresholdTotal<G, K, R> for Arranged<G, K, (), R, T1>
where
    G::Timestamp: TotalOrder+Lattice+Ord,
    T1: TraceReader<K, (), G::Timestamp, R>+Clone+'static,
    T1::Batch: BatchReader<K, (), G::Timestamp, R> {

    fn threshold_total<R2: Diff, F:Fn(R)->R2+'static>(&self, thresh: F) -> Collection<G, K, R2> {

        let mut trace = self.trace.clone();

        self.stream.unary_stream(Pipeline, "ThresholdTotal", move |input, output| {

            let thresh = &thresh;

            input.for_each(|capability, batches| {

                let mut session = output.session(&capability);
                for batch in batches.drain(..).map(|x| x.item) {

                    let mut batch_cursor = batch.cursor();
                    let (mut trace_cursor, trace_storage) = trace.cursor_through(batch.lower()).unwrap();

                    while batch_cursor.key_valid(&batch) {
                        let key = batch_cursor.key(&batch);
                        let mut count = R::zero();

                        // Compute the multiplicity of this key before the current batch.
                        trace_cursor.seek_key(&trace_storage, key);
                        if trace_cursor.key_valid(&trace_storage) && trace_cursor.key(&trace_storage) == key {
                            trace_cursor.map_times(&trace_storage, |_, diff| count = count + diff);
                        }

                        // Apply `thresh` both before and after `diff` is applied to `count`.
                        // If the result is non-zero, send it along.
                        batch_cursor.map_times(&batch, |time, diff| {
                            let old_weight = thresh(count);
                            count = count + diff;
                            let new_weight = thresh(count);
                            let difference = new_weight - old_weight;
                            if !difference.is_zero() {
                                session.give((key.clone(), time.clone(), difference));
                            }
                        });

                        batch_cursor.step_key(&batch);
                    }

                    // Tidy up the shared input trace.
                    trace.advance_by(batch.upper());
                    trace.distinguish_since(batch.upper());
                }
            });
        })
        .as_collection()
    }
}<commit_msg>correct bug for zero weights<commit_after>//! Reduce the collection to one occurrence of each distinct element.
//!
//! The `distinct_total` and `distinct_total_u` operators are optimizations of the more general
//! `distinct` and `distinct_u` operators for the case in which time is totally ordered.

use timely::order::TotalOrder;
use timely::dataflow::*;
use timely::dataflow::operators::Unary;
use timely::dataflow::channels::pact::Pipeline;

use lattice::Lattice;
use ::{Data, Collection, Diff};
use hashable::Hashable;
use collection::AsCollection;
use operators::arrange::{Arranged, ArrangeBySelf};
use trace::{BatchReader, Cursor, TraceReader};

/// Extension trait for the `distinct` differential dataflow method.
pub trait ThresholdTotal<G: Scope, K: Data, R: Diff> where G::Timestamp: TotalOrder+Lattice+Ord {
    /// Reduces the collection to one occurrence of each distinct element.
    ///
    /// # Examples
    ///
    /// ```
    /// extern crate timely;
    /// extern crate differential_dataflow;
    ///
    /// use differential_dataflow::input::Input;
    /// use differential_dataflow::operators::ThresholdTotal;
    ///
    /// fn main() {
    ///     ::timely::example(|scope| {
    ///         // report the number of occurrences of each key
    ///         scope.new_collection_from(1 .. 10).1
    ///              .map(|x| x / 3)
    ///              .threshold_total(|c| c % 2);
    ///     });
    /// }
    /// ```
    fn threshold_total<R2: Diff, F: Fn(R)->R2+'static>(&self, thresh: F) -> Collection<G, K, R2>;
    /// Reduces the collection to one occurrence of each distinct element.
    ///
    /// This reduction only tests whether the weight associated with a record is non-zero, and otherwise
    /// ignores its specific value. To take more general actions based on the accumulated weight, consider
    /// the `threshold` method.
    ///
    /// # Examples
    ///
    /// ```
    /// extern crate timely;
    /// extern crate differential_dataflow;
    ///
    /// use differential_dataflow::input::Input;
    /// use differential_dataflow::operators::ThresholdTotal;
    ///
    /// fn main() {
    ///     ::timely::example(|scope| {
    ///         // report the number of occurrences of each key
    ///         scope.new_collection_from(1 .. 10).1
    ///              .map(|x| x / 3)
    ///              .distinct_total();
    ///     });
    /// }
    /// ```
    fn distinct_total(&self) -> Collection<G, K, isize> {
        self.threshold_total(|c| if c.is_zero() { 0 } else { 1 })
    }
}

impl<G: Scope, K: Data+Hashable, R: Diff> ThresholdTotal<G, K, R> for Collection<G, K, R>
where G::Timestamp: TotalOrder+Lattice+Ord {
    fn threshold_total<R2: Diff, F: Fn(R)->R2+'static>(&self, thresh: F) -> Collection<G, K, R2> {
        self.arrange_by_self()
            .threshold_total(thresh)
    }
}

impl<G: Scope, K: Data, R: Diff, T1> ThresholdTotal<G, K, R> for Arranged<G, K, (), R, T1>
where
    G::Timestamp: TotalOrder+Lattice+Ord,
    T1: TraceReader<K, (), G::Timestamp, R>+Clone+'static,
    T1::Batch: BatchReader<K, (), G::Timestamp, R> {

    fn threshold_total<R2: Diff, F:Fn(R)->R2+'static>(&self, thresh: F) -> Collection<G, K, R2> {

        let mut trace = self.trace.clone();

        self.stream.unary_stream(Pipeline, "ThresholdTotal", move |input, output| {

            let thresh = &thresh;

            input.for_each(|capability, batches| {

                let mut session = output.session(&capability);
                for batch in batches.drain(..).map(|x| x.item) {

                    let mut batch_cursor = batch.cursor();
                    let (mut trace_cursor, trace_storage) = trace.cursor_through(batch.lower()).unwrap();

                    while batch_cursor.key_valid(&batch) {
                        let key = batch_cursor.key(&batch);
                        let mut count = R::zero();

                        // Compute the multiplicity of this key before the current batch.
                        trace_cursor.seek_key(&trace_storage, key);
                        if trace_cursor.key_valid(&trace_storage) && trace_cursor.key(&trace_storage) == key {
                            trace_cursor.map_times(&trace_storage, |_, diff| count = count + diff);
                        }

                        // Apply `thresh` both before and after `diff` is applied to `count`.
                        // If the result is non-zero, send it along.
                        batch_cursor.map_times(&batch, |time, diff| {
                            let old_weight = if count.is_zero() { R2::zero() } else { thresh(count) };
                            count = count + diff;
                            let new_weight = if count.is_zero() { R2::zero() } else { thresh(count) };
                            let difference = new_weight - old_weight;
                            if !difference.is_zero() {
                                session.give((key.clone(), time.clone(), difference));
                            }
                        });

                        batch_cursor.step_key(&batch);
                    }

                    // Tidy up the shared input trace.
                    trace.advance_by(batch.upper());
                    trace.distinguish_since(batch.upper());
                }
            });
        })
        .as_collection()
    }
}<|endoftext|>


<commit_before>use std::ops::{Mul, Add};

use nalgebra::{DMat, Transpose};
use rand::{Rand, Rng, Closed01};
use num::{Float, Zero};

use helpers::random01;

use ndarray::ArrayBase;
use ndarray::blas::AsBlas;

pub struct OrthogonalNMFBlas<FloatT> {
    pub hidden: ArrayBase<Vec<FloatT>, (usize, usize)>,
    pub weights: ArrayBase<Vec<FloatT>, (usize, usize)>,

    // these hold temporary results during an iteration.
    // kept in struct to prevent unnecessary memory allocations.
    // pub weights_dividend: ArrayBase<Vec<FloatT>, (usize, usize)>,
    // pub weights_divisor: ArrayBase<Vec<FloatT>, (usize, usize)>,
    // pub weights_divisor_reconstruction: ArrayBase<Vec<FloatT>, (usize, usize)>,
    //
    // pub hidden_dividend: ArrayBase<Vec<FloatT>, (usize, usize)>,
    // pub hidden_divisor: ArrayBase<Vec<FloatT>, (usize, usize)>,
    // pub hidden_divisor_tmp: ArrayBase<Vec<FloatT>, (usize, usize)>,
}

impl<FloatT> OrthogonalNMFBlas<FloatT>
    where FloatT: Float + Mul + Zero,
          Closed01<FloatT>: Rand
{
    pub fn init_random01<R: Rng>(nhidden: usize, nobserved: usize, nsamples: usize, rng: &mut R) -> OrthogonalNMFBlas<FloatT> {
        let mut hidden = ArrayBase::zeros((nhidden, nobserved));
        for x in hidden.iter_mut() {
            *x = random01(rng);
        }

        let mut weights = ArrayBase::zeros((nsamples, nhidden));
        for x in weights.iter_mut() {
            *x = random01(rng);
        }

        Self::init(hidden, weights)
    }

    // TODO initialize with values from last time step t-1 instead of choosing randomly
    pub fn init(hidden: ArrayBase<Vec<FloatT>, (usize, usize)>, weights: ArrayBase<Vec<FloatT>, (usize, usize)>) -> OrthogonalNMFBlas<FloatT> {
        OrthogonalNMFBlas {
            hidden: hidden,
            weights: weights,
        }
    }

    /// returns the number of observed variables
    #[inline]
    pub fn nobserved(&self) -> usize {
        self.hidden.shape()[1]
    }

    /// returns the number of hidden variables
    #[inline]
    pub fn nhidden(&self) -> usize {
        self.hidden.shape()[0]
    }

    /// returns the number of data points
    #[inline]
    pub fn nsamples(&self) -> usize {
        self.weights.shape()[0]
    }

    // TODO how many iterations ?
    // TODO compare this to the seoung solution
    /// it gets better and better with each iteration.
    /// one observed per column.
    /// one sample per row.
    pub fn iterate(&mut self, alpha: FloatT, data: &ArrayBase<Vec<FloatT>, (usize, usize)>) {
        assert_eq!(self.nsamples(), data.shape()[0]);
        assert_eq!(self.nobserved(), data.shape()[1]);
//
//         // has the same shape as weights
//         let new_weights_dividend = data.clone().mul(&hidden_transposed);
//         // has the same shape as weights
//         let new_weights_divisor = self.weights.clone().mul(&self.hidden).mul(&hidden_transposed);
//
//         // has the same shape as hidden
//         let new_hidden_dividend = weights_transposed.clone().mul(data);
//
//         // gamma is a symetric matrix with diagonal elements equal to zero
//         // and other elements = alpha
//         let gamma_size = self.nhidden();
//         let mut gamma = DMat::from_elem(gamma_size, gamma_size, alpha);
//
//         // set diagonal to zero
//         for i in 0..gamma_size {
//             gamma[(i, i)] = FloatT::zero();
//         }
//
//         // has the same shape as hidden
//         let new_hidden_divisor = weights_transposed.clone()
//             .mul(&self.weights)
//             .mul(&self.hidden)
//             // we add the previous latents
//             // multiplied by alpha except for the diag which is set to zero
//             .add(gamma.mul(&self.hidden));
//
//         // compute new weights
//         // TODO possibly use simd for this
//         for col in 0..self.weights.ncols() {
//             for row in 0..self.weights.nrows() {
//                 let index = (row, col);
//                 // if we have any zero in any of the matrizes
//                 // self.weights or self.hidden then
//                 // new_weights_divisor[index] will be zero
//                 // TODO how to deal with this
//                 let mut divisor = new_weights_divisor[index];
//                 if FloatT::zero() == divisor {
//                     divisor = FloatT::min_positive_value();
//                 }
//                 debug_assert!(FloatT::zero() != divisor);
//                 self.weights[index] =
//                     self.weights[index] *
//                     new_weights_dividend[index] /
//                     divisor;
//                 if FloatT::zero() == self.weights[index] {
//                     self.weights[index] = FloatT::min_positive_value();
//                 }
//                 debug_assert!(FloatT::zero() != self.weights[index]);
//             }
//         }
//
//         // compute new hidden
//         // TODO possibly use simd for this
//         for col in 0..self.hidden.ncols() {
//             for row in 0..self.hidden.nrows() {
//                 let index = (row, col);
//                 let mut divisor = new_hidden_divisor[index];
//                 if FloatT::zero() == divisor {
//                     divisor = FloatT::min_positive_value();
//                 }
//                 debug_assert!(FloatT::zero() != divisor);
//                 self.hidden[index] =
//                     self.hidden[index] *
//                     new_hidden_dividend[index] /
//                     divisor;
//                 if FloatT::zero() == self.hidden[index] {
//                     self.hidden[index] = FloatT::min_positive_value();
//                 }
//                 debug_assert!(FloatT::zero() != self.hidden[index]);
//             }
//         }
    }
}
<commit_msg>ortho nmf blas: flesh out<commit_after>use std::ops::{Mul, Add};

use rand::{Rand, Rng, Closed01};
use num::{Float, Zero};

use rblas::{Gemm, Matrix};
use rblas::attribute::Transpose;

use ndarray::ArrayBase;
use ndarray::blas::AsBlas;
use ndarray::blas;

use helpers::random01;

pub trait ShapeAsTuple<T> {
    fn shape_as_tuple(&self) -> T;
}

impl<T> ShapeAsTuple<(usize, usize)> for ArrayBase<Vec<T>, (usize, usize)> {
    #[inline]
    fn shape_as_tuple(&self) -> (usize, usize) {
        (self.shape()[0], self.shape()[1])
    }
}

// TODO Dimension type alias
// TODO OwnedArray

pub struct OrthogonalNMFBlas<FloatT> {
    pub hidden: ArrayBase<Vec<FloatT>, (usize, usize)>,
    pub weights: ArrayBase<Vec<FloatT>, (usize, usize)>,

    // these hold temporary results during an iteration.
    // kept in struct to prevent unnecessary memory allocations.
    pub tmp_weights_dividend: ArrayBase<Vec<FloatT>, (usize, usize)>,
    // pub tmp_weights_divisor: ArrayBase<Vec<FloatT>, (usize, usize)>,
    // pub tmp_weights_divisor_reconstruction: ArrayBase<Vec<FloatT>, (usize, usize)>,
    //
    // pub tmp_hidden_dividend: ArrayBase<Vec<FloatT>, (usize, usize)>,
    // pub tmp_hidden_divisor: ArrayBase<Vec<FloatT>, (usize, usize)>,
    // pub tmp_hidden_divisor_tmp: ArrayBase<Vec<FloatT>, (usize, usize)>,
}

impl<FloatT> OrthogonalNMFBlas<FloatT>
    where FloatT: Float + Mul + Zero + Clone + Gemm,
          // ArrayBase<Vec<FloatT>, (usize, usize)>: Matrix,
          // ndarray::blas::BlasArrayViewMut
          Closed01<FloatT>: Rand
{
    pub fn init_random01<R: Rng>(nhidden: usize, nobserved: usize, nsamples: usize, rng: &mut R) -> OrthogonalNMFBlas<FloatT> {
        let mut hidden = ArrayBase::zeros((nhidden, nobserved));
        for x in hidden.iter_mut() {
            *x = random01(rng);
        }

        let mut weights = ArrayBase::zeros((nsamples, nhidden));
        for x in weights.iter_mut() {
            *x = random01(rng);
        }

        Self::init(hidden, weights)
    }

    // TODO initialize with values from last time step t-1 instead of choosing randomly
    pub fn init(hidden: ArrayBase<Vec<FloatT>, (usize, usize)>, weights: ArrayBase<Vec<FloatT>, (usize, usize)>) -> OrthogonalNMFBlas<FloatT> {
        let weights_shape = weights.shape_as_tuple();
        OrthogonalNMFBlas {
            hidden: hidden,
            weights: weights,
            tmp_weights_dividend: ArrayBase::zeros(weights_shape),
        }
    }

    /// returns the number of observed variables
    #[inline]
    pub fn nobserved(&self) -> usize {
        self.hidden.shape()[1]
    }

    /// returns the number of hidden variables
    #[inline]
    pub fn nhidden(&self) -> usize {
        self.hidden.shape()[0]
    }

    /// returns the number of data points
    #[inline]
    pub fn nsamples(&self) -> usize {
        self.weights.shape()[0]
    }

    // TODO samples_shape ?
}

impl OrthogonalNMFBlas<f64> {
    // TODO how many iterations ?
    // TODO compare this to the seoung solution
    /// it gets better and better with each iteration.
    /// one observed per column.
    /// one sample per row.
    pub fn iterate(&mut self, alpha: f64, samples: &mut ArrayBase<Vec<f64>, (usize, usize)>) {
        assert_eq!(samples.shape_as_tuple(), (self.nsamples(), self.nobserved()));

        Gemm::gemm(
            &1.,
            Transpose::NoTrans, &samples.blas(),
            Transpose::Trans, &self.hidden.blas(),
            &0.,
            &mut self.tmp_weights_dividend.blas());
//
//         // has the same shape as weights
//         let new_weights_dividend = data.clone().mul(&hidden_transposed);
//         // has the same shape as weights
//         let new_weights_divisor = self.weights.clone().mul(&self.hidden).mul(&hidden_transposed);
//
//         // has the same shape as hidden
//         let new_hidden_dividend = weights_transposed.clone().mul(data);
//
//         // gamma is a symetric matrix with diagonal elements equal to zero
//         // and other elements = alpha
//         let gamma_size = self.nhidden();
//         let mut gamma = DMat::from_elem(gamma_size, gamma_size, alpha);
//
//         // set diagonal to zero
//         for i in 0..gamma_size {
//             gamma[(i, i)] = FloatT::zero();
//         }
//
//         // has the same shape as hidden
//         let new_hidden_divisor = weights_transposed.clone()
//             .mul(&self.weights)
//             .mul(&self.hidden)
//             // we add the previous latents
//             // multiplied by alpha except for the diag which is set to zero
//             .add(gamma.mul(&self.hidden));
//
//         // compute new weights
//         // TODO possibly use simd for this
//         for col in 0..self.weights.ncols() {
//             for row in 0..self.weights.nrows() {
//                 let index = (row, col);
//                 // if we have any zero in any of the matrizes
//                 // self.weights or self.hidden then
//                 // new_weights_divisor[index] will be zero
//                 // TODO how to deal with this
//                 let mut divisor = new_weights_divisor[index];
//                 if FloatT::zero() == divisor {
//                     divisor = FloatT::min_positive_value();
//                 }
//                 debug_assert!(FloatT::zero() != divisor);
//                 self.weights[index] =
//                     self.weights[index] *
//                     new_weights_dividend[index] /
//                     divisor;
//                 if FloatT::zero() == self.weights[index] {
//                     self.weights[index] = FloatT::min_positive_value();
//                 }
//                 debug_assert!(FloatT::zero() != self.weights[index]);
//             }
//         }
//
//         // compute new hidden
//         // TODO possibly use simd for this
//         for col in 0..self.hidden.ncols() {
//             for row in 0..self.hidden.nrows() {
//                 let index = (row, col);
//                 let mut divisor = new_hidden_divisor[index];
//                 if FloatT::zero() == divisor {
//                     divisor = FloatT::min_positive_value();
//                 }
//                 debug_assert!(FloatT::zero() != divisor);
//                 self.hidden[index] =
//                     self.hidden[index] *
//                     new_hidden_dividend[index] /
//                     divisor;
//                 if FloatT::zero() == self.hidden[index] {
//                     self.hidden[index] = FloatT::min_positive_value();
//                 }
//                 debug_assert!(FloatT::zero() != self.hidden[index]);
//             }
//         }
    }
}
<|endoftext|>


<commit_before>use std::env;
use std::io::prelude::*;
use std::fs::File;
use std::path::Path;

extern crate toml;

#[macro_use]
extern crate serde_derive;

// error in rust compiler.  Bugfix requested in Sept. 2017
// these are used, but the compiler is not seeing it for
// some reason
#[allow(unused_imports)]
#[macro_use]
extern crate serde_json;
// error in rust compiler.  Bugfix has been submitted in Sept. 2017
#[allow(unused_imports)]
#[macro_use]
extern crate serde;

// used in formatting the Cargo.toml file
#[derive(Deserialize, Debug)]
struct Tomlfile {
    contents: Contents,
}

// used in formatting the Cargo.toml file
#[derive(Deserialize, Debug)]
struct Metadata {
    deb: Deb,
}

// used in formatting the Cargo.toml file
#[derive(Deserialize, Debug)]
struct Deb {
    revision: Option<String>,
}

// used in formatting the Cargo.toml file
#[derive(Deserialize, Debug)]
struct Package {
    name: String,
    version: Option<String>,
    metadata: Metadata,
}

// used in formatting the Cargo.toml file
#[derive(Deserialize, Debug)]
struct Contents {
    package: Package,
    dependencies: Option<toml::Value>,
}

fn main() {
    let target = env::var("TARGET").unwrap();
    println!("target={}", target);

    if let Ok(_mode) = env::var("LIBINDY_STATIC") {
        let libindy_lib_path = env::var("LIBINDY_DIR").unwrap();
        println!("cargo:rustc-link-search=native={}",libindy_lib_path);
        println!("cargo:rustc-link-lib=static=indy");
    }else if target.contains("aarch64-linux-android") || target.contains("armv7-linux-androideabi") ||
        target.contains("arm-linux-androideabi") || target.contains("i686-linux-android") ||
        target.contains("x86_64-linux-android") || target.contains("aarch64-apple-ios") ||
        target.contains("armv7-apple-ios") || target.contains("armv7s-apple-ios") ||
        target.contains("i386-apple-ios") || target.contains("x86_64-apple-ios") {

        // let libindy_lib_path = match env::var("LIBINDY_DIR"){
        //     Ok(val) => val,
        //     Err(..) => panic!("Missing required environment variable LIBINDY_DIR")
        // };

        let openssl = match env::var("OPENSSL_LIB_DIR") {
            Ok(val) => val,
            Err(..) => match env::var("OPENSSL_DIR") {
                Ok(dir) => Path::new(&dir[..]).join("/lib").to_string_lossy().into_owned(),
                Err(..) => panic!("Missing required environment variables OPENSSL_DIR or OPENSSL_LIB_DIR")
            }
        };

        //println!("cargo:rustc-link-search=native={}",libindy_lib_path);
        //println!("cargo:rustc-link-lib=static=indy");
        println!("cargo:rustc-link-search=native={}", openssl);
        println!("cargo:rustc-link-lib=static=crypto");
        println!("cargo:rustc-link-lib=static=ssl");
        if cfg!(feature = "nullpay") {
            let libnullpay_lib_path = match env::var("LIBNULLPAY_DIR") {
                Ok(val) => val,
                Err(..) => panic!("Missing required environment variable LIBNULLPAY_DIR")
            };

            println!("cargo:rustc-link-search=native={}",libnullpay_lib_path);
            println!("cargo:rustc-link-lib=static=nullpay");
        } else if cfg!(feature = "sovtoken") {
            let libsovtoken_lib_path = match env::var("LIBSOVTOKEN_DIR") {
                Ok(val) => val,
                Err(..) => panic!("Missing required environment variable LIBSOVTOKEN_DIR")
            };

            println!("cargo:rustc-link-search=native={}",libsovtoken_lib_path);
            println!("cargo:rustc-link-lib=static=sovtoken");
        }
    }else if target.contains("darwin"){
        //OSX specific logic
        println!("cargo:rustc-link-lib=indy");
        //OSX does not allow 3rd party libs to be installed in /usr/lib. Instead install it in /usr/local/lib
        println!("cargo:rustc-link-search=native=/usr/local/lib");
        if cfg!(feature = "nullpay") {
          println!("cargo:rustc-link-lib=nullpay");
        } else if cfg!(feature = "sovtoken") {
            println!("cargo:rustc-link-lib=sovtoken");
        }
    }else if target.contains("-linux-"){
        //Linux specific logic
        println!("cargo:rustc-link-lib=indy");
        println!("cargo:rustc-link-search=native=/usr/lib");
        if cfg!(feature = "nullpay") {
          println!("cargo:rustc-link-lib=nullpay");
        } else if cfg!(feature = "sovtoken") {
            println!("cargo:rustc-link-lib=sovtoken");
        }
    }

    match env::var("CARGO_FEATURE_CI") {
        Ok(_) => {
            println!("injecting version information");
            // Leaving as unwrap, this is in the build script.
            let revision = get_revision().unwrap();
            write_variables(&revision);
        },
        Err(_) => {println!("NOT injecting version information"); },
    };
}

// Writes to the file 'src/utils/version_constants.rs' for use
// in outputing the version dynamically.
fn write_variables(revision:&str) {
    let out_dir = "src/utils/";
    let dest_path = Path::new(&out_dir).join("version_constants.rs");
    let mut f = File::create(&dest_path).unwrap();
    let s = format!("pub const VERSION: &'static str = env!(\"CARGO_PKG_VERSION\");\npub const REVISION: &'static str = \"{}\";\n", revision);
    if let Err(e) = f.write_all(s.as_bytes()) {
       panic!("Error creating version_constants.rs: {}", e);
    };

}

// Gets the revision number from the Cargo.toml file.
pub fn get_revision() -> Option<String> {
    let dir = match  env::var("CARGO_MANIFEST_DIR"){
        Ok(d) => d,
        Err(_) => panic!("Couldn't Manifest Directory"),
    };
    let filename = "Cargo.toml";
    let p = format!("{}/{}",dir,filename);
    let mut input = String::new();
    File::open(p).and_then(|mut f| {
        f.read_to_string(&mut input)}).unwrap();
    let tomlfile:Contents = toml::from_str(&input).unwrap();
    let revision:String = match tomlfile.package.metadata.deb.revision {
        Some(v) => v,
        None => String::from(""),
    };
    Some(format!("+{}", revision))
}

<commit_msg>Put back the LIBINDY_DIR variable.<commit_after>use std::env;
use std::io::prelude::*;
use std::fs::File;
use std::path::Path;

extern crate toml;

#[macro_use]
extern crate serde_derive;

// error in rust compiler.  Bugfix requested in Sept. 2017
// these are used, but the compiler is not seeing it for
// some reason
#[allow(unused_imports)]
#[macro_use]
extern crate serde_json;
// error in rust compiler.  Bugfix has been submitted in Sept. 2017
#[allow(unused_imports)]
#[macro_use]
extern crate serde;

// used in formatting the Cargo.toml file
#[derive(Deserialize, Debug)]
struct Tomlfile {
    contents: Contents,
}

// used in formatting the Cargo.toml file
#[derive(Deserialize, Debug)]
struct Metadata {
    deb: Deb,
}

// used in formatting the Cargo.toml file
#[derive(Deserialize, Debug)]
struct Deb {
    revision: Option<String>,
}

// used in formatting the Cargo.toml file
#[derive(Deserialize, Debug)]
struct Package {
    name: String,
    version: Option<String>,
    metadata: Metadata,
}

// used in formatting the Cargo.toml file
#[derive(Deserialize, Debug)]
struct Contents {
    package: Package,
    dependencies: Option<toml::Value>,
}

fn main() {
    let target = env::var("TARGET").unwrap();
    println!("target={}", target);

    if let Ok(_mode) = env::var("LIBINDY_STATIC") {
        let libindy_lib_path = env::var("LIBINDY_DIR").unwrap();
        println!("cargo:rustc-link-search=native={}",libindy_lib_path);
        println!("cargo:rustc-link-lib=static=indy");
    }else if target.contains("aarch64-linux-android") || target.contains("armv7-linux-androideabi") ||
        target.contains("arm-linux-androideabi") || target.contains("i686-linux-android") ||
        target.contains("x86_64-linux-android") || target.contains("aarch64-apple-ios") ||
        target.contains("armv7-apple-ios") || target.contains("armv7s-apple-ios") ||
        target.contains("i386-apple-ios") || target.contains("x86_64-apple-ios") {

        let libindy_lib_path = match env::var("LIBINDY_DIR"){
            Ok(val) => val,
            Err(..) => panic!("Missing required environment variable LIBINDY_DIR")
        };

        let openssl = match env::var("OPENSSL_LIB_DIR") {
            Ok(val) => val,
            Err(..) => match env::var("OPENSSL_DIR") {
                Ok(dir) => Path::new(&dir[..]).join("/lib").to_string_lossy().into_owned(),
                Err(..) => panic!("Missing required environment variables OPENSSL_DIR or OPENSSL_LIB_DIR")
            }
        };

        println!("cargo:rustc-link-search=native={}",libindy_lib_path);
        println!("cargo:rustc-link-lib=static=indy");
        println!("cargo:rustc-link-search=native={}", openssl);
        println!("cargo:rustc-link-lib=static=crypto");
        println!("cargo:rustc-link-lib=static=ssl");
        if cfg!(feature = "nullpay") {
            let libnullpay_lib_path = match env::var("LIBNULLPAY_DIR") {
                Ok(val) => val,
                Err(..) => panic!("Missing required environment variable LIBNULLPAY_DIR")
            };

            println!("cargo:rustc-link-search=native={}",libnullpay_lib_path);
            println!("cargo:rustc-link-lib=static=nullpay");
        } else if cfg!(feature = "sovtoken") {
            let libsovtoken_lib_path = match env::var("LIBSOVTOKEN_DIR") {
                Ok(val) => val,
                Err(..) => panic!("Missing required environment variable LIBSOVTOKEN_DIR")
            };

            println!("cargo:rustc-link-search=native={}",libsovtoken_lib_path);
            println!("cargo:rustc-link-lib=static=sovtoken");
        }
    }else if target.contains("darwin"){
        //OSX specific logic
        println!("cargo:rustc-link-lib=indy");
        //OSX does not allow 3rd party libs to be installed in /usr/lib. Instead install it in /usr/local/lib
        println!("cargo:rustc-link-search=native=/usr/local/lib");
        if cfg!(feature = "nullpay") {
          println!("cargo:rustc-link-lib=nullpay");
        } else if cfg!(feature = "sovtoken") {
            println!("cargo:rustc-link-lib=sovtoken");
        }
    }else if target.contains("-linux-"){
        //Linux specific logic
        println!("cargo:rustc-link-lib=indy");
        println!("cargo:rustc-link-search=native=/usr/lib");
        if cfg!(feature = "nullpay") {
          println!("cargo:rustc-link-lib=nullpay");
        } else if cfg!(feature = "sovtoken") {
            println!("cargo:rustc-link-lib=sovtoken");
        }
    }

    match env::var("CARGO_FEATURE_CI") {
        Ok(_) => {
            println!("injecting version information");
            // Leaving as unwrap, this is in the build script.
            let revision = get_revision().unwrap();
            write_variables(&revision);
        },
        Err(_) => {println!("NOT injecting version information"); },
    };
}

// Writes to the file 'src/utils/version_constants.rs' for use
// in outputing the version dynamically.
fn write_variables(revision:&str) {
    let out_dir = "src/utils/";
    let dest_path = Path::new(&out_dir).join("version_constants.rs");
    let mut f = File::create(&dest_path).unwrap();
    let s = format!("pub const VERSION: &'static str = env!(\"CARGO_PKG_VERSION\");\npub const REVISION: &'static str = \"{}\";\n", revision);
    if let Err(e) = f.write_all(s.as_bytes()) {
       panic!("Error creating version_constants.rs: {}", e);
    };

}

// Gets the revision number from the Cargo.toml file.
pub fn get_revision() -> Option<String> {
    let dir = match  env::var("CARGO_MANIFEST_DIR"){
        Ok(d) => d,
        Err(_) => panic!("Couldn't Manifest Directory"),
    };
    let filename = "Cargo.toml";
    let p = format!("{}/{}",dir,filename);
    let mut input = String::new();
    File::open(p).and_then(|mut f| {
        f.read_to_string(&mut input)}).unwrap();
    let tomlfile:Contents = toml::from_str(&input).unwrap();
    let revision:String = match tomlfile.package.metadata.deb.revision {
        Some(v) => v,
        None => String::from(""),
    };
    Some(format!("+{}", revision))
}

<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use dom::bindings::codegen::Bindings::WebSocketBinding;
use dom::bindings::codegen::Bindings::WebSocketBinding::WebSocketMethods;
use dom::bindings::codegen::Bindings::EventHandlerBinding::EventHandlerNonNull;
use dom::bindings::codegen::InheritTypes::EventTargetCast;
use dom::bindings::error::{Error, ErrorResult, Fallible};
use dom::bindings::error::Error::{InvalidState, InvalidAccess};
use dom::bindings::error::Error::{Network, Syntax, Security, Abort, Timeout};
use dom::event::{Event, EventBubbles, EventCancelable, EventHelpers};
use dom::bindings::global::{GlobalField, GlobalRef, GlobalRoot};
use dom::bindings::js::{Temporary, JSRef};
use dom::bindings::utils::reflect_dom_object;
use dom::eventtarget::{EventTarget, EventTargetHelpers, EventTargetTypeId};
use util::str::DOMString;

use websocket::{Message, Sender, Receiver};
use websocket::client::request::Url;
use websocket::Client;
use std::cell::Cell;
use std::borrow::ToOwned;

#[derive(PartialEq, Copy)]
#[jstraceable]
enum WebSocketRequestState {
	Connecting = 0,
    Open = 1,
    Closing = 2,
    Closed = 3,
}

#[dom_struct]
pub struct WebSocket {
    eventtarget: EventTarget,
    url: DOMString,
    global: GlobalField,
	ready_state: Cell<WebSocketRequestState>
}

impl WebSocket {
    pub fn new_inherited(global: GlobalRef, url: DOMString) -> WebSocket {
        println!("Creating websocket...");
	let copied_url = url.clone();
	WebSocket {
            eventtarget: EventTarget::new_inherited(EventTargetTypeId::WebSocket),
            url: url,
            global: GlobalField::from_rooted(&global),
		ready_state: Cell::new(WebSocketRequestState::Connecting)
        }

    }

    pub fn new(global: GlobalRef, url: DOMString) -> Temporary<WebSocket> {
        let ws_root = reflect_dom_object(box WebSocket::new_inherited(global, url),
                           global,
                           WebSocketBinding::Wrap).root();
        ws_root.r().Open();
        Temporary::from_rooted(ws_root.r())
    }

    pub fn Constructor(global: GlobalRef, url: DOMString) -> Fallible<Temporary<WebSocket>> {
    	Ok(WebSocket::new(global, url))
    }

/*	fn send(self) -> ErrorResult {
		tx_1 = self.tx.clone();
		let send_loop = thread::scoped(move || {
			loop {
				let message = match self.rx.recv() {
					Ok(m) => m,
					Err(e) => {
						println!("Send loop: {:?}",e);
						return;
					}
				};
				match message {
					Message::Close(_) => {
						let _ = self.sender.send_message(message);
						// If it's a close message, send it and return
						return;
					}
					_ => (),
				}
				//Send the message
				match self.sender.send_message(message){
					Ok(()) => (),
					Err(e) => {
						println!("Send Loop: {:?}", e);
						let _ = self.sender.send_message(Message::Close(None));
						return;
					}
				}
			}
		});
	}*/
}

impl<'a> WebSocketMethods for JSRef<'a, WebSocket> {
	event_handler!(open, GetOnopen, SetOnopen);
	event_handler!(close, GetOnclose, SetOnclose);
	event_handler!(error, GetOnerror, SetOnerror);

    fn Url(self) -> DOMString {
	    println!("Cloning URL");
      self.url.clone()
    }

   fn ReadyState(self) -> u16 {
   	println!("Setting readystate");
	  self.ready_state.get() as u16
   }

   fn Open(self) -> ErrorResult {
    	println!("Trying to connect.");
		let parsed_url = Url::parse(self.url.as_slice()).unwrap();
   		let request = Client::connect(parsed_url).unwrap();
		let response = request.send().unwrap();
		response.validate().unwrap();
		println!("Successful connection.");
    	let global = self.global.root();
        let event = Event::new(global.r(),
                               "open".to_owned(),
                               EventBubbles::DoesNotBubble,
                               EventCancelable::Cancelable).root();
        let target: JSRef<EventTarget> = EventTargetCast::from_ref(self);
        event.r().fire(target);
        println!("Fired event.");
		Ok(())
    }
}
<commit_msg>Open is now asynchronous, events now working correctly.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use dom::bindings::codegen::Bindings::WebSocketBinding;
use dom::bindings::codegen::Bindings::WebSocketBinding::WebSocketMethods;
use dom::bindings::codegen::Bindings::EventHandlerBinding::EventHandlerNonNull;
use dom::bindings::codegen::InheritTypes::EventTargetCast;
use dom::bindings::error::{Error, ErrorResult, Fallible};
use dom::bindings::error::Error::{InvalidState, InvalidAccess};
use dom::bindings::error::Error::{Network, Syntax, Security, Abort, Timeout};
use dom::event::{Event, EventBubbles, EventCancelable, EventHelpers};
use dom::bindings::global::{GlobalField, GlobalRef, GlobalRoot};
use dom::bindings::js::{Temporary, JSRef};
use dom::bindings::utils::reflect_dom_object;
use dom::eventtarget::{EventTarget, EventTargetHelpers, EventTargetTypeId};
use util::str::DOMString;
use script_task::Runnable;
use script_task::ScriptMsg;
use dom::bindings::refcounted::Trusted;

use websocket::{Message, Sender, Receiver};
use websocket::client::request::Url;
use websocket::Client;
use std::cell::Cell;
use std::borrow::ToOwned;

#[derive(PartialEq, Copy)]
#[jstraceable]
enum WebSocketRequestState {
	Connecting = 0,
    Open = 1,
    Closing = 2,
    Closed = 3,
}

#[dom_struct]
pub struct WebSocket {
    eventtarget: EventTarget,
    url: DOMString,
    global: GlobalField,
	ready_state: Cell<WebSocketRequestState>
}

impl WebSocket {
    pub fn new_inherited(global: GlobalRef, url: DOMString) -> WebSocket {
        println!("Creating websocket...");
	let copied_url = url.clone();
	WebSocket {
            eventtarget: EventTarget::new_inherited(EventTargetTypeId::WebSocket),
            url: url,
            global: GlobalField::from_rooted(&global),
		ready_state: Cell::new(WebSocketRequestState::Connecting)
        }

    }

    pub fn new(global: GlobalRef, url: DOMString) -> Temporary<WebSocket> {
        let ws_root = reflect_dom_object(box WebSocket::new_inherited(global, url),
                           global,
                           WebSocketBinding::Wrap).root();
        ws_root.r().Open();
        Temporary::from_rooted(ws_root.r())
    }

    pub fn Constructor(global: GlobalRef, url: DOMString) -> Fallible<Temporary<WebSocket>> {
    	Ok(WebSocket::new(global, url))
    }

/*	fn send(self) -> ErrorResult {
		tx_1 = self.tx.clone();
		let send_loop = thread::scoped(move || {
			loop {
				let message = match self.rx.recv() {
					Ok(m) => m,
					Err(e) => {
						println!("Send loop: {:?}",e);
						return;
					}
				};
				match message {
					Message::Close(_) => {
						let _ = self.sender.send_message(message);
						// If it's a close message, send it and return
						return;
					}
					_ => (),
				}
				//Send the message
				match self.sender.send_message(message){
					Ok(()) => (),
					Err(e) => {
						println!("Send Loop: {:?}", e);
						let _ = self.sender.send_message(Message::Close(None));
						return;
					}
				}
			}
		});
	}*/

    fn Close(self, code: Option<u16>, reason: Option<DOMString>) -> Fallible<()>{
	if(code.is_some()){ //Code defined
	   if(!(code==Some(1000)||
		(code>=Some(3000) && code<=Some(4999))
	       )
	   )
	   { //Check code is NOT 1000 NOR in the range of 3000-4999 (inclusive)
		return Err(Error::InvalidAccess); //Throw InvalidAccessError and abort
	   }
	}
	if(reason.is_some()){ //reason defined
	   //FIX ME
	   /*if(reason.capacity() > 123) //reason cannot be larger than 123 bytes
	   {
		return Err(Error::Syntax); //Throw SyntaxError and abort
	   }*/
	}
	//TODO:
	/*match self.ready_state.get() { //Returns the value of the cell
	   //WebsocketRequestState::Closing => (), //Do nothing
	   //WebsocketRequestState::Closed => (), //Do nothing
	   //To do:
	   //How to detect not yet established - Receiving state?
	      //Fail the WebSocket connection - how? What does this really mean for the websocket object?
	      //Set readyState to closing
	   //
	   //How to detect not yet been started - Unsent state?
	      //Start the Websocket closing handshake - how? What does this really mean for the websocket object?
	      //if code.is_some - WebSocket status code in close message to be the same as code
	      //if reason.is_some - Websocket close message reason to be same as reason
	   //_ => {self.ready_state.set(WebsocketRequestState::Closing);}
	}*/
	Ok(())
    }
}

impl<'a> WebSocketMethods for JSRef<'a, WebSocket> {
	event_handler!(open, GetOnopen, SetOnopen);
	event_handler!(close, GetOnclose, SetOnclose);
	event_handler!(error, GetOnerror, SetOnerror);

    fn Url(self) -> DOMString {
	    println!("Cloning URL");
      self.url.clone()
    }

   fn ReadyState(self) -> u16 {
   	println!("Setting readystate");
	  self.ready_state.get() as u16
   }

   fn Open(self) -> ErrorResult {
   		let global_root = self.global.root();
        let addr: Trusted<WebSocket> = Trusted::new(global_root.r().get_cx(), self, global_root.r().script_chan().clone());
    	let open_task = box WebSocketTaskHandler::new(addr.clone(), WebSocketTask::Open);
        global_root.r().script_chan().send(ScriptMsg::RunnableMsg(open_task)).unwrap();
   		Ok(())
    }

    fn Close(self, code: Option<u16>, reason: Option<DOMString>) -> Fallible<()>{
	if(code.is_some()){ //Code defined
	   if(!(code==Some(1000)||
		(code>=Some(3000) && code<=Some(4999))
	       )
	   )
	   { //Check code is NOT 1000 NOR in the range of 3000-4999 (inclusive)
		return Err(Error::InvalidAccess); //Throw InvalidAccessError and abort
	   }
	}
	if(reason.is_some()){ //reason defined
	   //FIX ME
	   /*if(reason.capacity() > 123) //reason cannot be larger than 123 bytes
	   {
		return Err(Error::Syntax); //Throw SyntaxError and abort
	   }*/
	}
	//TODO:
	/*match self.ready_state.get() { //Returns the value of the cell
	   //WebsocketRequestState::Closing => (), //Do nothing
	   //WebsocketRequestState::Closed => (), //Do nothing
	   //To do:
	   //How to detect not yet established - Receiving state?
	      //Fail the WebSocket connection - how? What does this really mean for the websocket object?
	      //Set readyState to closing
	   //
	   //How to detect not yet been started - Unsent state?
	      //Start the Websocket closing handshake - how? What does this really mean for the websocket object?
	      //if code.is_some - WebSocket status code in close message to be the same as code
	      //if reason.is_some - Websocket close message reason to be same as reason
	   //_ => {self.ready_state.set(WebsocketRequestState::Closing);}
	}*/
	Ok(())
    }
}

pub enum WebSocketTask {
    Open,
}

pub struct WebSocketTaskHandler {
    addr: Trusted<WebSocket>,
    task: WebSocketTask,
}

impl WebSocketTaskHandler {
    pub fn new(addr: Trusted<WebSocket>, task: WebSocketTask) -> WebSocketTaskHandler {
        WebSocketTaskHandler {
            addr: addr,
            task: task,
        }
    }

    fn dispatch_open(&self) {
    	println!("Trying to connect.");
    	let ws = self.addr.to_temporary().root();
		let parsed_url = Url::parse(ws.r().url.as_slice()).unwrap();
   		let request = Client::connect(parsed_url).unwrap();
		let response = request.send().unwrap();
		response.validate().unwrap();
		println!("Successful connection.");
    	let global = ws.r().global.root();
        let event = Event::new(global.r(),
                               "open".to_owned(),
                               EventBubbles::DoesNotBubble,
                               EventCancelable::Cancelable).root();
        let target: JSRef<EventTarget> = EventTargetCast::from_ref(ws.r());
        event.r().fire(target);
        println!("Fired event.");
    }
}

impl Runnable for WebSocketTaskHandler {
    fn handler(self: Box<WebSocketTaskHandler>) {
        match self.task {
            WebSocketTask::Open => {
                self.dispatch_open();
            }
        }
    }
}

<|endoftext|>