<commit_before>use std::thread;
use std::path::Path;

use websocket::{Server, Message, Sender, Receiver};
use websocket::message::Type;
use websocket::header::WebSocketProtocol;

use git2::Repository;
use git2::ObjectType;

use payload;
use State;

extern crate fsevent;
use std::sync::mpsc::channel;

use serde_xml as xml;

pub fn start() {
    let server = Server::bind("127.0.0.1:2794").unwrap();
    thread::spawn(move || {
        for connection in server {
            thread::spawn(move || {
                let request = connection.unwrap().read_request().unwrap();
                let headers = request.headers.clone();
                request.validate().unwrap();
                let mut response = request.accept();
                println!("Connected");

                if let Some(&WebSocketProtocol(ref protocols)) = headers.get() {
                    if protocols.contains(&("rust-websocket".to_string())) {
                        response.headers.set(WebSocketProtocol(vec!["superconductor".to_string()]));
                    }
                }

                let mut client = response.send().unwrap();
                let message = Message::text(payload::generate(None));
                let (tx, rx) = channel();
                client.send_message(&message).unwrap();

                let monitor = thread::spawn(move || {
                    println!("Monitoring");
                    let fsevent = fsevent::FsEvent::new(tx);
                    fsevent.append_path(".");
                    fsevent.observe();
                });

                let (mut sender, mut receiver) = client.split();

                let notifier = thread::spawn(move || {
                    loop {
                        let value = rx.recv().unwrap();
                        println!("{:?}", value);
                        let message = Message::text(payload::generate(None));
                        sender.send_message(&message).unwrap();
                        //let message = Message::text(payload::generate(Some(state)));
                        //sender.send_message(&message).unwrap();
                    }
                });

                for message in receiver.incoming_messages() {
                    let message: Message = message.unwrap_or(Message::close());
                    match message.opcode {
                        Type::Close => {
                            let message = Message::close();
                            //sender.send_message(&message).unwrap();
                            monitor.join();
                            println!("Monitor stopped");
                            notifier.join();
                            println!("Notifier stopped");
                            println!("Client disconnected");
                            return;
                        },
                        Type::Ping => {
                            let message = Message::pong(message.payload);
                            //sender.send_message(&message).unwrap();
                        },
                        _ => {
                            let payload = String::from_utf8_lossy(message.payload.as_ref());
                            println!("{:?}", payload);
                            let mut state: State = xml::from_str(&payload).unwrap();
                            println!("{:?}", state);

                            let repo = Repository::discover(".").unwrap();
                            let head = repo.head().unwrap().peel(ObjectType::Commit).unwrap();
                            let mut index = repo.index().unwrap();

                            let to_remove = index.iter().fold(vec![], |mut acc, entry| {
                                let entry_path = String::from_utf8_lossy(entry.path.as_ref());
                                match state.include.iter().find(|i| i.as_ref() == entry_path) {
                                    None => acc.push(entry_path.into_owned()),
                                    _ => {}
                                };
                                acc
                            });

                            repo.reset_default(Some(&head), to_remove.iter()).unwrap();
                            for change in state.clone().include {
                                let path = Path::new(&change);
                                index.add_path(path).unwrap();
                            }
                            index.write().unwrap();

                            if let Some(event) = state.save_update.clone() {
                                let author = repo.signature().unwrap();
                                index.read(false);
                                let tree_oid = index.write_tree().unwrap();
                                let tree = repo.find_tree(tree_oid).unwrap();
                                repo.commit(Some("HEAD"), &author, &author, &state.message, &tree, &[&head.as_commit().unwrap()]);
                            }
                        }
                    }
                }
            });
        }
    });
}

<commit_msg>Working refactor with single receiver<commit_after>use std::thread;
use std::thread::JoinHandle;
use std::path::Path;

use websocket::{WebSocketStream, Server};
use websocket::Message as WebMessage;
use websocket::Sender as WebSender;
use websocket::client::Sender as WebClientSender;
use websocket::Receiver as WebReceiver;
use websocket::message::Type as WebMessageType;
use websocket::server::Connection;
use websocket::header::WebSocketProtocol;

use git2::Repository;
use git2::ObjectType;

use payload;
use State;

extern crate fsevent;
use self::fsevent::Event as FsEvent;
use std::sync::mpsc::channel;
use std::sync::mpsc::Sender;
use std::sync::mpsc::Receiver;

use serde_xml as xml;

#[derive(Debug)]
enum NotifierMessage<'a> {
    WebMessage(WebMessage<'a>),
    FsEvent(FsEvent)
}

pub fn start() {
    let server = Server::bind("127.0.0.1:2794").unwrap();
    thread::spawn(move || {
        for connection in server {
            thread::spawn(move || connect(connection.unwrap()));
        }
    });
}

fn connect(connection: Connection<WebSocketStream, WebSocketStream>) {
    let request = connection.read_request().unwrap();
    let headers = request.headers.clone();
    request.validate().unwrap();
    let mut response = request.accept();
    println!("Connected");

    if let Some(&WebSocketProtocol(ref protocols)) = headers.get() {
        if protocols.contains(&("rust-websocket".to_string())) {
            response.headers.set(WebSocketProtocol(vec!["superconductor".to_string()]));
        }
    }

    let mut client = response.send().unwrap();
    let message = WebMessage::text(payload::generate(None));
    client.send_message(&message).unwrap();

    let (tx, rx) = channel::<NotifierMessage>();
    let (mut sender, mut receiver) = client.split();
    thread::spawn(move || start_notifier(rx, sender));
    let txc = tx.clone();
    thread::spawn(move || start_monitor(txc));
    for message in receiver.incoming_messages() {
        let message: WebMessage = message.unwrap_or(WebMessage::close());
        tx.send(NotifierMessage::WebMessage(message));
    }
}

fn start_notifier(rx: Receiver<NotifierMessage>, mut sender: WebClientSender<WebSocketStream>) {
    loop {
        let value = rx.recv().unwrap();
        println!("{:?}", value);
        match value {
            NotifierMessage::WebMessage(message) => {
                match message.opcode {
                    WebMessageType::Close => {
                        let message = WebMessage::close();
                        //sender.send_message(&message).unwrap();
                        println!("Client disconnected");
                        return;
                    },
                    WebMessageType::Ping => {
                        let message = WebMessage::pong(message.payload);
                        //sender.send_message(&message).unwrap();
                    },
                    _ => {
                        let payload = String::from_utf8_lossy(message.payload.as_ref());
                        println!("{:?}", payload);
                        let mut state: State = xml::from_str(&payload).unwrap();
                        println!("{:?}", state);

                        let repo = Repository::discover(".").unwrap();
                        let head = repo.head().unwrap().peel(ObjectType::Commit).unwrap();
                        let mut index = repo.index().unwrap();

                        let to_remove = index.iter().fold(vec![], |mut acc, entry| {
                            let entry_path = String::from_utf8_lossy(entry.path.as_ref());
                            match state.include.iter().find(|i| i.as_ref() == entry_path) {
                                None => acc.push(entry_path.into_owned()),
                                _ => {}
                            };
                            acc
                        });

                        repo.reset_default(Some(&head), to_remove.iter()).unwrap();
                        for change in state.clone().include {
                            let path = Path::new(&change);
                            index.add_path(path).unwrap();
                        }
                        index.write().unwrap();

                        if let Some(event) = state.save_update.clone() {
                            let author = repo.signature().unwrap();
                            index.read(false);
                            let tree_oid = index.write_tree().unwrap();
                            let tree = repo.find_tree(tree_oid).unwrap();
                            repo.commit(Some("HEAD"), &author, &author, &state.message, &tree, &[&head.as_commit().unwrap()]);
                            let message = WebMessage::text(payload::generate(None));
                            sender.send_message(&message).unwrap();
                        } else {
                            let message = WebMessage::text(payload::generate(Some(state)));
                            sender.send_message(&message).unwrap();
                        }
                    }
                }
            },
            NotifierMessage::FsEvent(event) => {}
        }
    }
}

fn start_monitor(tx: Sender<NotifierMessage>) {
    let (ftx, rx) = channel();
    thread::spawn(move || {
        println!("Monitoring");
        let fsevent = fsevent::FsEvent::new(ftx);
        fsevent.append_path(".");
        fsevent.observe();
    });
    loop {
        let event = rx.recv().unwrap();
        tx.send(NotifierMessage::FsEvent(event));
    }
}

//for message in receiver.incoming_messages() {

//}
<|endoftext|>


<commit_before>extern crate dbus;
use self::dbus::{NameFlag, ReleaseNameReply};

use super::connection::DBusConnection;
use super::error::DBusError;
use super::interface::DBusInterfaceMap;
use super::message::{DBusMessage, DBusMessageType};
use super::object::DBusObject;
use super::target::DBusTarget;

use std::collections::btree_map::{BTreeMap, Entry};

pub type DBusSignalHandler = Box<FnMut(&DBusConnection, &DBusTarget) -> ()>;
type DBusSignalHandlers = Vec<DBusSignalHandler>;
type DBusSignalHandlerMap = BTreeMap<DBusTarget, DBusSignalHandlers>;

fn _add_handler(handlers: &mut DBusSignalHandlerMap, signal: DBusTarget, handler: DBusSignalHandler) {
    match handlers.entry(signal) {
        Entry::Vacant(v)    => { v.insert(vec![handler]); },
        Entry::Occupied(o)  => o.into_mut().push(handler),
    };
}

pub struct DBusServer<'a> {
    conn: &'a DBusConnection,
    name: String,

    objects: BTreeMap<String, DBusObject<'a>>,
    signals: DBusSignalHandlerMap,
    namespace_signals: DBusSignalHandlerMap,
}

impl<'a> DBusServer<'a> {
    pub fn new(conn: &'a DBusConnection, name: &str) -> Result<DBusServer<'a>, DBusError> {
        try!(conn.register_name(name, NameFlag::DoNotQueue as u32));

        Ok(DBusServer {
            conn: conn,
            name: name.to_owned(),

            objects: BTreeMap::new(),
            signals: BTreeMap::new(),
            namespace_signals: BTreeMap::new(),
        })
    }

    pub fn name(&self) -> &str {
        &self.name
    }

    pub fn add_object(&mut self, path: &str, iface_map: DBusInterfaceMap<'a>) -> Result<&mut Self, DBusError> {
        match self.objects.entry(path.to_owned()) {
            Entry::Vacant(v)    => {
                let obj = try!(DBusObject::new(self.conn, path, iface_map));

                v.insert(obj);

                Ok(())
            },
            Entry::Occupied(_)  => Err(DBusError::PathAlreadyRegistered(path.to_owned())),
        }.map(|_| self)
    }

    pub fn remove_object(&mut self, path: &str) -> Result<&mut Self, DBusError> {
        match self.objects.remove(path) {
            Some(_) => Ok(self),
            None    => Err(DBusError::NoSuchPath(path.to_owned())),
        }
    }

    pub fn connect(&mut self, signal: DBusTarget, callback: DBusSignalHandler) -> Result<&mut Self, DBusError> {
        try!(self.conn.add_match(&format!("type='signal',interface='{}',path='{}',member='{}'",
                                          signal.interface,
                                          signal.object,
                                          signal.method)));

        _add_handler(&mut self.signals, signal, callback);

        Ok(self)
    }

    pub fn connect_namespace(&mut self, signal: DBusTarget, callback: DBusSignalHandler) -> Result<&mut Self, DBusError> {
        try!(self.conn.add_match(&format!("type='signal',interface='{}',path_namespace='{}',member='{}'",
                                          signal.interface,
                                          signal.object,
                                          signal.method)));

        _add_handler(&mut self.namespace_signals, signal, callback);

        Ok(self)
    }

    pub fn handle_message<'b>(&mut self, m: &'b mut DBusMessage) -> Option<&'b mut DBusMessage> {
        match m.msg_type() {
            DBusMessageType::Signal     => Some(self._match_signal(m)),
            DBusMessageType::MethodCall => self._call_method(m),
            _                           => Some(m),
        }
    }

    fn _call_method<'b>(&mut self, m: &'b mut DBusMessage) -> Option<&'b mut DBusMessage> {
        self.objects.iter_mut().fold(Some(m), |opt_m, (_, object)| {
            opt_m.and_then(|mut m| {
                match object.handle_message(&mut m) {
                    None          => Some(m),
                    Some(Ok(()))  => None,
                    Some(Err(())) => {
                        println!("failed to send a reply for {:?}", m);
                        None
                    },
                }
            })
        })
    }

    fn _match_signal<'b>(&mut self, m: &'b mut DBusMessage) -> &'b mut DBusMessage {
        let conn = (&self.conn).clone();

        DBusTarget::extract(m).map(|signal| {
            self.signals.get_mut(&signal).map(|handlers| {
                handlers.iter_mut().map(|f| {
                    f(conn, &signal);
                })
            });

            self.namespace_signals.iter_mut().filter(|&(expect, _)| {
                expect.namespace_eq(&signal)
            }).map(|(_, handlers)| {
                handlers.iter_mut().map(|f| {
                    f(conn, &signal);
                })
            }).collect::<Vec<_>>();
        });

        m
    }
}

impl<'a> Drop for DBusServer<'a> {
    fn drop(&mut self) {
        let res = self.conn.release_name(&self.name);
        match res {
            Ok(reply) =>
                match reply {
                    ReleaseNameReply::Released    => (),
                    ReleaseNameReply::NonExistent => panic!("internal error: non-existent name {}?!", self.name),
                    ReleaseNameReply::NotOwner    => panic!("internal error: not the owner of {}?!", self.name),
                },
            Err(err) => println!("failed to release {}: {:?}: {:?}", self.name, err.name(), err.message()),
        }
    }
}
<commit_msg>server: use for loops in signal handling<commit_after>extern crate dbus;
use self::dbus::{NameFlag, ReleaseNameReply};

use super::connection::DBusConnection;
use super::error::DBusError;
use super::interface::DBusInterfaceMap;
use super::message::{DBusMessage, DBusMessageType};
use super::object::DBusObject;
use super::target::DBusTarget;

use std::collections::btree_map::{BTreeMap, Entry};

pub type DBusSignalHandler = Box<FnMut(&DBusConnection, &DBusTarget) -> ()>;
type DBusSignalHandlers = Vec<DBusSignalHandler>;
type DBusSignalHandlerMap = BTreeMap<DBusTarget, DBusSignalHandlers>;

fn _add_handler(handlers: &mut DBusSignalHandlerMap, signal: DBusTarget, handler: DBusSignalHandler) {
    match handlers.entry(signal) {
        Entry::Vacant(v)    => { v.insert(vec![handler]); },
        Entry::Occupied(o)  => o.into_mut().push(handler),
    };
}

pub struct DBusServer<'a> {
    conn: &'a DBusConnection,
    name: String,

    objects: BTreeMap<String, DBusObject<'a>>,
    signals: DBusSignalHandlerMap,
    namespace_signals: DBusSignalHandlerMap,
}

impl<'a> DBusServer<'a> {
    pub fn new(conn: &'a DBusConnection, name: &str) -> Result<DBusServer<'a>, DBusError> {
        try!(conn.register_name(name, NameFlag::DoNotQueue as u32));

        Ok(DBusServer {
            conn: conn,
            name: name.to_owned(),

            objects: BTreeMap::new(),
            signals: BTreeMap::new(),
            namespace_signals: BTreeMap::new(),
        })
    }

    pub fn name(&self) -> &str {
        &self.name
    }

    pub fn add_object(&mut self, path: &str, iface_map: DBusInterfaceMap<'a>) -> Result<&mut Self, DBusError> {
        match self.objects.entry(path.to_owned()) {
            Entry::Vacant(v)    => {
                let obj = try!(DBusObject::new(self.conn, path, iface_map));

                v.insert(obj);

                Ok(())
            },
            Entry::Occupied(_)  => Err(DBusError::PathAlreadyRegistered(path.to_owned())),
        }.map(|_| self)
    }

    pub fn remove_object(&mut self, path: &str) -> Result<&mut Self, DBusError> {
        match self.objects.remove(path) {
            Some(_) => Ok(self),
            None    => Err(DBusError::NoSuchPath(path.to_owned())),
        }
    }

    pub fn connect(&mut self, signal: DBusTarget, callback: DBusSignalHandler) -> Result<&mut Self, DBusError> {
        try!(self.conn.add_match(&format!("type='signal',interface='{}',path='{}',member='{}'",
                                          signal.interface,
                                          signal.object,
                                          signal.method)));

        _add_handler(&mut self.signals, signal, callback);

        Ok(self)
    }

    pub fn connect_namespace(&mut self, signal: DBusTarget, callback: DBusSignalHandler) -> Result<&mut Self, DBusError> {
        try!(self.conn.add_match(&format!("type='signal',interface='{}',path_namespace='{}',member='{}'",
                                          signal.interface,
                                          signal.object,
                                          signal.method)));

        _add_handler(&mut self.namespace_signals, signal, callback);

        Ok(self)
    }

    pub fn handle_message<'b>(&mut self, m: &'b mut DBusMessage) -> Option<&'b mut DBusMessage> {
        match m.msg_type() {
            DBusMessageType::Signal     => Some(self._match_signal(m)),
            DBusMessageType::MethodCall => self._call_method(m),
            _                           => Some(m),
        }
    }

    fn _call_method<'b>(&mut self, m: &'b mut DBusMessage) -> Option<&'b mut DBusMessage> {
        self.objects.iter_mut().fold(Some(m), |opt_m, (_, object)| {
            opt_m.and_then(|mut m| {
                match object.handle_message(&mut m) {
                    None          => Some(m),
                    Some(Ok(()))  => None,
                    Some(Err(())) => {
                        println!("failed to send a reply for {:?}", m);
                        None
                    },
                }
            })
        })
    }

    fn _match_signal<'b>(&mut self, m: &'b mut DBusMessage) -> &'b mut DBusMessage {
        let conn = (&self.conn).clone();

        DBusTarget::extract(m).map(|signal| {
            for handlers in self.signals.get_mut(&signal) {
                for handler in handlers.iter_mut() {
                    handler(conn, &signal);
                }
            }

            let matched_handlers = self.namespace_signals.iter_mut().filter(|&(expect, _)| {
                expect.namespace_eq(&signal)
            });

            for (_, handlers) in matched_handlers {
                for handler in handlers.iter_mut() {
                    handler(conn, &signal);
                };
            };
        });

        m
    }
}

impl<'a> Drop for DBusServer<'a> {
    fn drop(&mut self) {
        let res = self.conn.release_name(&self.name);
        match res {
            Ok(reply) =>
                match reply {
                    ReleaseNameReply::Released    => (),
                    ReleaseNameReply::NonExistent => panic!("internal error: non-existent name {}?!", self.name),
                    ReleaseNameReply::NotOwner    => panic!("internal error: not the owner of {}?!", self.name),
                },
            Err(err) => println!("failed to release {}: {:?}: {:?}", self.name, err.name(), err.message()),
        }
    }
}
<|endoftext|>


<commit_before>use std::marker::PhantomData;
use std::iter;

pub trait Shrink {
    type Item;
    type Iterator: Iterator<Item=Self::Item>;

    fn shrink(value: &Self::Item) -> Self::Iterator;
}

pub struct Empty<T>(PhantomData<T>);

impl <T> Empty<T> {
    pub fn empty() -> Self { Empty(PhantomData) }
}

impl <T> Shrink for Empty<T> {
    type Item = T;
    type Iterator = iter::Empty<T>;

    fn shrink(_: &Self::Item) -> Self::Iterator {
        iter::empty()
    }
}

macro_rules! tuple_impls {
    ($($name:ident),*) => {
        impl <$($name: 'static + Shrink),*> Shrink for ($($name,)*) {
            type Item = ($($name::Item,)*);
            type Iterator = Box<Iterator<Item=Self::Item>>;

            #[inline]
            #[allow(unused_variables, non_snake_case)]
            fn shrink(value: &Self::Item) -> Self::Iterator {
                let ( $(ref $name,)* ) = *value;
                Box::new(
                    iter::empty()
                )
//                ($($name.generate(ctx),)*)
            }
        }
    }
}

macro_tuples_impl!{tuple_impls}

pub struct DefaultShrinker<T>(PhantomData<T>);

impl <T> DefaultShrinker<T>
    where DefaultShrinker<T>: Shrink
{
    pub fn new() -> Self { DefaultShrinker(PhantomData) }
}

impl Shrink for DefaultShrinker<bool> {
    type Item = bool;
    type Iterator = Box<Iterator<Item=bool>>;

    fn shrink(value: &bool) -> Self::Iterator {
        match *value {
            true => Box::new(iter::once(false)),
            false => Box::new(iter::empty())
        }
    }
}
<commit_msg>Shrink impls for tuples.<commit_after>use std::marker::PhantomData;
use std::iter;

pub trait Shrink {
    type Item;
    type Iterator: Iterator<Item=Self::Item>;

    fn shrink(value: &Self::Item) -> Self::Iterator;
}

pub struct Empty<T>(PhantomData<T>);

impl <T> Empty<T> {
    pub fn empty() -> Self { Empty(PhantomData) }
}

impl <T> Shrink for Empty<T> {
    type Item = T;
    type Iterator = iter::Empty<T>;

    fn shrink(_: &Self::Item) -> Self::Iterator {
        iter::empty()
    }
}

macro_rules! tuple_impls {
    ($first:ident $(, $rest:ident)*) => {
        impl <$first: 'static + Shrink $(, $rest: 'static + Shrink)*> Shrink for ($first, $($rest),*)
            where $first::Item: Clone $(, $rest::Item: Clone)*
        {
            type Item = ($first::Item, $($rest::Item),*);
            type Iterator = Box<Iterator<Item=Self::Item>>;

            #[inline]
            #[allow(unused_variables, non_snake_case)]
            fn shrink(value: &Self::Item) -> Self::Iterator {
                let ( ref $first, $(ref $rest),*) = *value;
                let t_rest: ($($rest::Item,)*) = ($($rest.clone(),)*);

                Box::new(
                    <$first as Shrink>::shrink($first)
                        .scan(t_rest.clone(), |rest, $first| {
                            let ($($rest,)*) = rest.clone();
                            Some(($first, $($rest),*))
                        })
                        .chain(
                            <($($rest,)*) as Shrink>::shrink(&t_rest)
                                .scan($first.clone(), |$first, ($($rest,)*)| {
                                    Some(($first.clone(), $($rest),*))
                                })
                        )
                )
            }
        }
    };
    () => {
        impl Shrink for () {
            type Item = ();
            type Iterator = iter::Empty<()>;

            #[inline]
            fn shrink(_: &()) -> Self::Iterator {
                iter::empty()
            }
        }
    }
}

macro_tuples_impl!{tuple_impls}

pub struct DefaultShrinker<T>(PhantomData<T>);

impl <T> DefaultShrinker<T>
    where DefaultShrinker<T>: Shrink
{
    pub fn new() -> Self { DefaultShrinker(PhantomData) }
}

impl Shrink for DefaultShrinker<bool> {
    type Item = bool;
    type Iterator = Box<Iterator<Item=bool>>;

    fn shrink(value: &bool) -> Self::Iterator {
        match *value {
            true => Box::new(iter::once(false)),
            false => Box::new(iter::empty())
        }
    }
}
<|endoftext|>


<commit_before>#![crate_name = "sknife"]
use std::collections::HashMap;

/// Finds list of all keys in a Hashmap satisfying a predicate
///
/// # Arguments
///
/// * `predicate` - The predicate function
/// * `obj` - A hashmap/object to find the keys in
///
/// # Example
///
/// ```
/// use sknife::find_keys::find_keys;
/// use std::collections::HashMap;
/// let obj = HashMap::new();
/// let greater_than_one = |x: i32| x > 1;
/// find_keys(greater_than_one, obj);
///
/// ```
///
/// # Result
/// ```
/// vec![String::from("Three"), String::from("Four")];
/// ```
pub fn find_keys<V, F> (mut predicate: F, obj: HashMap<String, V>) -> Vec<String>
    where F: FnMut(V) -> bool,
    V: Clone {
        let mut list: Vec<String> = vec![];
        for (k, v) in obj {
            match predicate(v.clone()) {
                true => {
                    list.push(k);
                },
                false => continue
            }
        }

        list
}

mod tests {
    use super::*;
    use std::collections::HashMap;

    #[test]
    fn find_keys_hash() {
        let mut obj = HashMap::new();
        obj.insert(String::from("One"), 1);
        obj.insert(String::from("Three"), 3);
        obj.insert(String::from("Four"), 4);
        let greater_than_one = |x: i32| x > 1;
        assert_eq!(
            find_keys(greater_than_one, obj),
            vec![
                String::from("Four"),
                String::from("Three")
            ]
        );
    }

    #[test]
    fn find_keys_empty_hash() {
        let mut obj = HashMap::new();
        let greater_than_one = |x: i32| x > 1;
        let expected_result: Vec<String> = vec![];
        assert_eq!(
            find_keys(greater_than_one, obj),
            expected_result
        );
    }
}<commit_msg>Add the order of keys correctly<commit_after>#![crate_name = "sknife"]
use std::collections::HashMap;

/// Finds list of all keys in a Hashmap satisfying a predicate
///
/// # Arguments
///
/// * `predicate` - The predicate function
/// * `obj` - A hashmap/object to find the keys in
///
/// # Example
///
/// ```
/// use sknife::find_keys::find_keys;
/// use std::collections::HashMap;
/// let obj = HashMap::new();
/// let greater_than_one = |x: i32| x > 1;
/// find_keys(greater_than_one, obj);
///
/// ```
///
/// # Result
/// ```
/// vec![String::from("Three"), String::from("Four")];
/// ```
pub fn find_keys<V, F> (mut predicate: F, obj: HashMap<String, V>) -> Vec<String>
    where F: FnMut(V) -> bool,
    V: Clone {
        let mut list: Vec<String> = vec![];
        for (k, v) in obj {
            match predicate(v.clone()) {
                true => {
                    list.push(k);
                },
                false => continue
            }
        }

        list
}

mod tests {
    use super::*;
    use std::collections::HashMap;

    #[test]
    fn find_keys_hash() {
        let mut obj = HashMap::new();
        obj.insert(String::from("One"), 1);
        obj.insert(String::from("Three"), 3);
        obj.insert(String::from("Four"), 4);
        let greater_than_one = |x: i32| x > 1;
        assert_eq!(
            find_keys(greater_than_one, obj),
            vec![
                String::from("Three"),
                String::from("Four")
            ]
        );
    }

    #[test]
    fn find_keys_empty_hash() {
        let mut obj = HashMap::new();
        let greater_than_one = |x: i32| x > 1;
        let expected_result: Vec<String> = vec![];
        assert_eq!(
            find_keys(greater_than_one, obj),
            expected_result
        );
    }
}<|endoftext|>


<commit_before>use std::io::{Read, Write, Result, Error};
use std::net::{ToSocketAddrs};
use socket::UtpSocket;

/// Stream interface for UtpSocket.
pub struct UtpStream {
    socket: UtpSocket,
}

impl UtpStream {
    /// Create a uTP stream listening on the given address.
    #[unstable]
    pub fn bind<A: ToSocketAddrs>(addr: A) -> Result<UtpStream> {
        match UtpSocket::bind(addr) {
            Ok(s)  => Ok(UtpStream { socket: s }),
            Err(e) => Err(e),
        }
    }

    /// Open a uTP connection to a remote host by hostname or IP address.
    #[unstable]
    pub fn connect<A: ToSocketAddrs>(dst: A) -> Result<UtpStream> {
        // Port 0 means the operating system gets to choose it
        let my_addr = "0.0.0.0:0";
        let socket = match UtpSocket::bind(my_addr) {
            Ok(s) => s,
            Err(e) => return Err(e),
        };

        match socket.connect(dst) {
            Ok(socket) => Ok(UtpStream { socket: socket }),
            Err(e) => Err(e),
        }
    }

    /// Gracefully close connection to peer.
    ///
    /// This method allows both peers to receive all packets still in
    /// flight.
    #[unstable]
    pub fn close(&mut self) -> Result<()> {
        self.socket.close()
    }
}

impl Read for UtpStream {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
        match self.socket.recv_from(buf) {
            Ok((read, _src)) => Ok(read),
            Err(e) => Err(e),
        }
    }
}

impl Write for UtpStream {
    fn write(&mut self, buf: &[u8]) -> Result<usize> {
        self.socket.send_to(buf)
    }

    // TODO: Actually implement flushing
    fn flush(&mut self) -> Result<()> {
        Ok(())
    }
}
<commit_msg>Remove unused import.<commit_after>use std::io::{Read, Write, Result};
use std::net::{ToSocketAddrs};
use socket::UtpSocket;

/// Stream interface for UtpSocket.
pub struct UtpStream {
    socket: UtpSocket,
}

impl UtpStream {
    /// Create a uTP stream listening on the given address.
    #[unstable]
    pub fn bind<A: ToSocketAddrs>(addr: A) -> Result<UtpStream> {
        match UtpSocket::bind(addr) {
            Ok(s)  => Ok(UtpStream { socket: s }),
            Err(e) => Err(e),
        }
    }

    /// Open a uTP connection to a remote host by hostname or IP address.
    #[unstable]
    pub fn connect<A: ToSocketAddrs>(dst: A) -> Result<UtpStream> {
        // Port 0 means the operating system gets to choose it
        let my_addr = "0.0.0.0:0";
        let socket = match UtpSocket::bind(my_addr) {
            Ok(s) => s,
            Err(e) => return Err(e),
        };

        match socket.connect(dst) {
            Ok(socket) => Ok(UtpStream { socket: socket }),
            Err(e) => Err(e),
        }
    }

    /// Gracefully close connection to peer.
    ///
    /// This method allows both peers to receive all packets still in
    /// flight.
    #[unstable]
    pub fn close(&mut self) -> Result<()> {
        self.socket.close()
    }
}

impl Read for UtpStream {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
        match self.socket.recv_from(buf) {
            Ok((read, _src)) => Ok(read),
            Err(e) => Err(e),
        }
    }
}

impl Write for UtpStream {
    fn write(&mut self, buf: &[u8]) -> Result<usize> {
        self.socket.send_to(buf)
    }

    // TODO: Actually implement flushing
    fn flush(&mut self) -> Result<()> {
        Ok(())
    }
}
<|endoftext|>


<commit_before>use std::cmp;
use std::collections::BinaryHeap;
use std::io;

use fst::build::Builder;
use fst::{Fst, FstReader, Output};
use Result;

pub fn union_ignore_outputs<'f, W, I>(
    bfst: &mut Builder<W>,
    fsts: I,
) -> Result<()>
where W: io::Write,
      I: IntoIterator<Item=&'f Fst> {
    let mut union = Union::new(fsts);
    while let Some(slot) = union.pop() {
        try!(bfst.add(slot.input()));
        union.refill(slot);
    }
    Ok(())
}

pub fn union_with_outputs<'f, W, F, I>(
    bfst: &mut Builder<W>,
    fsts: I,
    mut merge: F,
) -> Result<()>
where W: io::Write,
      F: FnMut(&[u8], &[(usize, u64)]) -> u64,
      I: IntoIterator<Item=&'f Fst> {
    let mut union = Union::new(fsts);
    let mut outs = vec![];
    while let Some(slot) = union.pop() {
        unsafe { outs.set_len(0); }
        outs.push((slot.index(), slot.output().value()));
        while union.peek_is_duplicate(slot.input()) {
            let slot2 = union.pop().unwrap();
            outs.push((slot2.index(), slot2.output().value()));
            union.refill(slot2);
        }
        try!(bfst.insert(slot.input(), merge(slot.input(), &outs)));
        union.refill(slot);
    }
    Ok(())
}

struct Union<'f> {
    rdrs: Vec<FstReader<'f>>,
    heap: BinaryHeap<Slot>,
}

impl<'f> Union<'f> {
    fn new<I: IntoIterator<Item=&'f Fst>>(fsts: I) -> Union<'f> {
        let mut u = Union {
            rdrs: fsts.into_iter().map(Fst::reader).collect(),
            heap: BinaryHeap::new(),
        };
        for i in 0..u.rdrs.len() {
            u.refill(Slot::new(i));
        }
        u
    }

    fn pop(&mut self) -> Option<Slot> {
        self.heap.pop()
    }

    fn peek_is_duplicate(&self, key: &[u8]) -> bool {
        self.heap.peek().map(|s| s.input() == key).unwrap_or(false)
    }

    fn refill(&mut self, mut slot: Slot) {
        if let Some((input, output)) = self.rdrs[slot.idx].next() {
            slot.set_input(input);
            slot.set_output(output);
            self.heap.push(slot);
        }
    }
}

#[derive(Debug, Eq, PartialEq)]
struct Slot {
    idx: usize,
    input: Vec<u8>,
    output: Output,
}

impl Slot {
    fn new(rdr_idx: usize) -> Slot {
        Slot {
            idx: rdr_idx,
            input: Vec::with_capacity(64),
            output: Output::zero(),
        }
    }

    fn index(&self) -> usize {
        self.idx
    }

    fn input(&self) -> &[u8] {
        &self.input
    }

    fn output(&self) -> Output {
        self.output
    }

    fn set_input(&mut self, input: &[u8]) {
        let addcap = input.len().checked_sub(self.input.len()).unwrap_or(0);
        self.input.reserve(addcap);
        unsafe { self.input.set_len(input.len()); }
        for (dst, &byte) in self.input.iter_mut().zip(input.iter()) {
            *dst = byte;
        }
    }

    fn set_output(&mut self, output: Output) {
        self.output = output;
    }
}

impl PartialOrd for Slot {
    fn partial_cmp(&self, other: &Slot) -> Option<cmp::Ordering> {
        (&self.input, self.output)
        .partial_cmp(&(&other.input, other.output))
        .map(|ord| ord.reverse())
    }
}

impl Ord for Slot {
    fn cmp(&self, other: &Slot) -> cmp::Ordering {
        self.partial_cmp(other).unwrap()
    }
}

#[cfg(test)]
mod tests {
    use fst::build::Builder;
    use fst::tests::{fst_map, fst_set, fst_inputstrs_outputs, fst_input_strs};
    use fst::Fst;
    use super::{union_ignore_outputs, union_with_outputs};

    fn s(string: &str) -> String { string.to_owned() }

    #[test]
    fn basic_set() {
        let set1 = fst_set(&["a", "b", "c"]);
        let set2 = fst_set(&["x", "y", "z"]);

        let mut bfst = Builder::memory();
        union_ignore_outputs(&mut bfst, &[set1, set2]).unwrap();
        let fst = Fst::from_bytes(bfst.into_inner().unwrap()).unwrap();
        assert_eq!(fst_input_strs(&fst), vec!["a", "b", "c", "x", "y", "z"]);
    }

    #[test]
    fn basic_set_dupes() {
        let set1 = fst_set(&["aa", "b", "cc"]);
        let set2 = fst_set(&["b", "cc", "z"]);

        let mut bfst = Builder::memory();
        union_ignore_outputs(&mut bfst, &[set1, set2]).unwrap();
        let fst = Fst::from_bytes(bfst.into_inner().unwrap()).unwrap();
        assert_eq!(fst_input_strs(&fst), vec!["aa", "b", "cc", "z"]);
    }

    #[test]
    fn basic_map() {
        let map1 = fst_map(vec![("a", 1), ("b", 2), ("c", 3)]);
        let map2 = fst_map(vec![("x", 1), ("y", 2), ("z", 3)]);

        let mut bfst = Builder::memory();
        union_with_outputs(
            &mut bfst,
            &[map1, map2],
            |_, os| os.iter().fold(0, |a, &(_, b)| a + b),
        ).unwrap();
        let fst = Fst::from_bytes(bfst.into_inner().unwrap()).unwrap();
        assert_eq!(
            fst_inputstrs_outputs(&fst),
            vec![
                (s("a"), 1), (s("b"), 2), (s("c"), 3),
                (s("x"), 1), (s("y"), 2), (s("z"), 3),
            ]);
    }

    #[test]
    fn basic_map_dupes() {
        let map1 = fst_map(vec![("aa", 1), ("b", 2), ("cc", 3)]);
        let map2 = fst_map(vec![("b", 1), ("cc", 2), ("z", 3)]);
        let map3 = fst_map(vec![("b", 1)]);

        let mut bfst = Builder::memory();
        union_with_outputs(
            &mut bfst,
            &[map1, map2, map3],
            |_, os| os.iter().fold(0, |a, &(_, b)| a + b),
        ).unwrap();
        let fst = Fst::from_bytes(bfst.into_inner().unwrap()).unwrap();
        assert_eq!(
            fst_inputstrs_outputs(&fst),
            vec![
                (s("aa"), 1), (s("b"), 4), (s("cc"), 5), (s("z"), 3),
            ]);
    }
}
<commit_msg>Give more information to the merge function.<commit_after>use std::cmp;
use std::collections::BinaryHeap;
use std::io;

use fst::build::Builder;
use fst::{Fst, FstReader, Output};
use Result;

pub fn union_ignore_outputs<'f, W, I>(
    bfst: &mut Builder<W>,
    fsts: I,
) -> Result<()>
where W: io::Write,
      I: IntoIterator<Item=&'f Fst> {
    let mut union = Union::new(fsts);
    while let Some(slot) = union.pop() {
        try!(bfst.add(slot.input()));
        union.refill(slot);
    }
    Ok(())
}

#[derive(Copy, Clone, Debug)]
pub struct FstOutput {
    pub index: usize,
    pub output: u64,
}

pub fn union_with_outputs<'f, W, F, I>(
    bfst: &mut Builder<W>,
    fsts: I,
    mut merge: F,
) -> Result<()>
where W: io::Write,
      F: FnMut(&[u8], &[FstOutput]) -> u64,
      I: IntoIterator<Item=&'f Fst> {
    let mut union = Union::new(fsts);
    let mut outs = vec![];
    while let Some(slot) = union.pop() {
        unsafe { outs.set_len(0); }
        outs.push(slot.fst_output());
        while union.peek_is_duplicate(slot.input()) {
            let slot2 = union.pop().unwrap();
            outs.push(slot2.fst_output());
            union.refill(slot2);
        }
        try!(bfst.insert(slot.input(), merge(slot.input(), &outs)));
        union.refill(slot);
    }
    Ok(())
}

struct Union<'f> {
    rdrs: Vec<FstReader<'f>>,
    heap: BinaryHeap<Slot>,
}

impl<'f> Union<'f> {
    fn new<I: IntoIterator<Item=&'f Fst>>(fsts: I) -> Union<'f> {
        let mut u = Union {
            rdrs: fsts.into_iter().map(Fst::reader).collect(),
            heap: BinaryHeap::new(),
        };
        for i in 0..u.rdrs.len() {
            u.refill(Slot::new(i));
        }
        u
    }

    fn pop(&mut self) -> Option<Slot> {
        self.heap.pop()
    }

    fn peek_is_duplicate(&self, key: &[u8]) -> bool {
        self.heap.peek().map(|s| s.input() == key).unwrap_or(false)
    }

    fn refill(&mut self, mut slot: Slot) {
        if let Some((input, output)) = self.rdrs[slot.idx].next() {
            slot.set_input(input);
            slot.set_output(output);
            self.heap.push(slot);
        }
    }
}

#[derive(Debug, Eq, PartialEq)]
struct Slot {
    idx: usize,
    input: Vec<u8>,
    output: Output,
}

impl Slot {
    fn new(rdr_idx: usize) -> Slot {
        Slot {
            idx: rdr_idx,
            input: Vec::with_capacity(64),
            output: Output::zero(),
        }
    }

    fn fst_output(&self) -> FstOutput {
        FstOutput { index: self.idx, output: self.output.value() }
    }

    fn input(&self) -> &[u8] {
        &self.input
    }

    fn output(&self) -> Output {
        self.output
    }

    fn set_input(&mut self, input: &[u8]) {
        let addcap = input.len().checked_sub(self.input.len()).unwrap_or(0);
        self.input.reserve(addcap);
        unsafe { self.input.set_len(input.len()); }
        for (dst, &byte) in self.input.iter_mut().zip(input.iter()) {
            *dst = byte;
        }
    }

    fn set_output(&mut self, output: Output) {
        self.output = output;
    }
}

impl PartialOrd for Slot {
    fn partial_cmp(&self, other: &Slot) -> Option<cmp::Ordering> {
        (&self.input, self.output)
        .partial_cmp(&(&other.input, other.output))
        .map(|ord| ord.reverse())
    }
}

impl Ord for Slot {
    fn cmp(&self, other: &Slot) -> cmp::Ordering {
        self.partial_cmp(other).unwrap()
    }
}

#[cfg(test)]
mod tests {
    use fst::build::Builder;
    use fst::tests::{fst_map, fst_set, fst_inputstrs_outputs, fst_input_strs};
    use fst::Fst;
    use super::{union_ignore_outputs, union_with_outputs};

    fn s(string: &str) -> String { string.to_owned() }

    #[test]
    fn basic_set() {
        let set1 = fst_set(&["a", "b", "c"]);
        let set2 = fst_set(&["x", "y", "z"]);

        let mut bfst = Builder::memory();
        union_ignore_outputs(&mut bfst, &[set1, set2]).unwrap();
        let fst = Fst::from_bytes(bfst.into_inner().unwrap()).unwrap();
        assert_eq!(fst_input_strs(&fst), vec!["a", "b", "c", "x", "y", "z"]);
    }

    #[test]
    fn basic_set_dupes() {
        let set1 = fst_set(&["aa", "b", "cc"]);
        let set2 = fst_set(&["b", "cc", "z"]);

        let mut bfst = Builder::memory();
        union_ignore_outputs(&mut bfst, &[set1, set2]).unwrap();
        let fst = Fst::from_bytes(bfst.into_inner().unwrap()).unwrap();
        assert_eq!(fst_input_strs(&fst), vec!["aa", "b", "cc", "z"]);
    }

    #[test]
    fn basic_map() {
        let map1 = fst_map(vec![("a", 1), ("b", 2), ("c", 3)]);
        let map2 = fst_map(vec![("x", 1), ("y", 2), ("z", 3)]);

        let mut bfst = Builder::memory();
        union_with_outputs(
            &mut bfst,
            &[map1, map2],
            |_, os| os.iter().fold(0, |a, b| a + b.output),
        ).unwrap();
        let fst = Fst::from_bytes(bfst.into_inner().unwrap()).unwrap();
        assert_eq!(
            fst_inputstrs_outputs(&fst),
            vec![
                (s("a"), 1), (s("b"), 2), (s("c"), 3),
                (s("x"), 1), (s("y"), 2), (s("z"), 3),
            ]);
    }

    #[test]
    fn basic_map_dupes() {
        let map1 = fst_map(vec![("aa", 1), ("b", 2), ("cc", 3)]);
        let map2 = fst_map(vec![("b", 1), ("cc", 2), ("z", 3)]);
        let map3 = fst_map(vec![("b", 1)]);

        let mut bfst = Builder::memory();
        union_with_outputs(
            &mut bfst,
            &[map1, map2, map3],
            |_, os| os.iter().fold(0, |a, b| a + b.output),
        ).unwrap();
        let fst = Fst::from_bytes(bfst.into_inner().unwrap()).unwrap();
        assert_eq!(
            fst_inputstrs_outputs(&fst),
            vec![
                (s("aa"), 1), (s("b"), 4), (s("cc"), 5), (s("z"), 3),
            ]);
    }
}
<|endoftext|>


<commit_before>// This file is part of libfringe, a low-level green threading library.
// Copyright (c) whitequark <whitequark@whitequark.org>
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

//! Generators.
//!
//! Generators allow repeatedly suspending the execution of a function,
//! returning a value to the caller, and resuming the suspended function
//! afterwards.

use core::marker::PhantomData;
use core::{ptr, mem};

use stack;
use context::Context;

#[derive(Debug, Clone, Copy)]
pub enum State {
  /// Generator can be resumed. This is the initial state.
  Runnable,
  /// Generator cannot be resumed. This is the state of the generator after
  /// the generator function has returned or panicked.
  Unavailable
}

/// Generator wraps a function and allows suspending its execution more than once, returning
/// a value each time.
///
/// The first time `resume(input0)` is called, the function is called as `f(yielder, input0)`.
/// It runs until it suspends its execution through `yielder.suspend(output0)`, after which
/// `resume(input0)` returns `output0`. The function can be resumed again using `resume(input1)`,
/// after which `yielder.suspend(output0)` returns `input1`, and so on. Once the function returns,
/// the `resume()` call will return `None`, and it will return `None` every time it is called
/// after that.
///
/// If the generator function panics, the panic is propagated through the `resume()` call as usual.
///
/// After the generator function returns or panics, it is safe to reclaim the generator stack
/// using `unwrap()`.
///
/// `state()` can be used to determine whether the generator function has returned;
/// the state is `State::Runnable` after creation and suspension, and `State::Unavailable`
/// once the generator function returns or panics.
///
/// When the input type is `()`, a generator implements the Iterator trait.
///
/// # Example
///
/// ```
/// use fringe::{OsStack, Generator};
///
/// let stack = OsStack::new(0).unwrap();
/// let mut add_one = Generator::new(stack, move |yielder, mut input| {
///   loop {
///     if input == 0 { break }
///     input = yielder.suspend(input + 1)
///   }
/// });
/// println!("{:?}", add_one.resume(2)); // prints Some(3)
/// println!("{:?}", add_one.resume(3)); // prints Some(4)
/// println!("{:?}", add_one.resume(0)); // prints None
/// ```
///
/// # Iterator example
///
/// ```
/// use fringe::{OsStack, Generator};
///
/// let stack = OsStack::new(0).unwrap();
/// let mut nat = Generator::new(stack, move |yielder, ()| {
///   for i in 1.. { yielder.suspend(i) }
/// });
/// println!("{:?}", nat.next()); // prints Some(0)
/// println!("{:?}", nat.next()); // prints Some(1)
/// println!("{:?}", nat.next()); // prints Some(2)
/// ```
#[derive(Debug)]
pub struct Generator<Input: Send, Output: Send, Stack: stack::Stack> {
  state:   State,
  context: Context<Stack>,
  phantom: (PhantomData<*const Input>, PhantomData<*const Output>)
}

impl<Input, Output, Stack> Generator<Input, Output, Stack>
    where Input: Send, Output: Send, Stack: stack::Stack {
  /// Creates a new generator.
  ///
  /// See also the [contract](../trait.GuardedStack.html) that needs to be fulfilled by `stack`.
  pub fn new<F>(stack: Stack, f: F) -> Generator<Input, Output, Stack>
      where Stack: stack::GuardedStack,
            F: FnOnce(&mut Yielder<Input, Output, Stack>, Input) + Send {
    unsafe { Generator::unsafe_new(stack, f) }
  }

  /// Same as `new`, but does not require `stack` to have a guard page.
  ///
  /// This function is unsafe because the generator function can easily violate
  /// memory safety by overflowing the stack. It is useful in environments where
  /// guarded stacks do not exist, e.g. in absence of an MMU.
  ///
  /// See also the [contract](../trait.Stack.html) that needs to be fulfilled by `stack`.
  pub unsafe fn unsafe_new<F>(stack: Stack, f: F) -> Generator<Input, Output, Stack>
      where F: FnOnce(&mut Yielder<Input, Output, Stack>, Input) + Send {
    unsafe extern "C" fn generator_wrapper<Input, Output, Stack, F>(env: usize) -> !
        where Input: Send, Output: Send, Stack: stack::Stack,
              F: FnOnce(&mut Yielder<Input, Output, Stack>, Input) {
      // Retrieve our environment from the callee and return control to it.
      let (mut yielder, f) = ptr::read(env as *mut (Yielder<Input, Output, Stack>, F));
      let data = Context::swap(yielder.context, yielder.context, 0);
      // See the second half of Yielder::suspend_bare.
      let (new_context, input) = ptr::read(data as *mut (*mut Context<Stack>, Input));
      yielder.context = new_context as *mut Context<Stack>;
      // Run the body of the generator.
      f(&mut yielder, input);
      // Past this point, the generator has dropped everything it has held.
      loop { yielder.suspend_bare(None); }
    }

    let mut generator = Generator {
      state:   State::Runnable,
      context: Context::new(stack, generator_wrapper::<Input, Output, Stack, F>),
      phantom: (PhantomData, PhantomData)
    };

    // Transfer environment to the callee.
    let mut env = (Yielder::new(&mut generator.context), f);
    Context::swap(&mut generator.context, &generator.context,
                  &mut env as *mut (Yielder<Input, Output, Stack>, F) as usize);
    mem::forget(env);

    generator
  }

  /// Resumes the generator and return the next value it yields.
  /// If the generator function has returned, returns `None`.
  #[inline]
  pub fn resume(&mut self, input: Input) -> Option<Output> {
    match self.state {
      State::Runnable => {
        // Set the state to Unavailable. Since we have exclusive access to the generator,
        // the only case where this matters is the generator function panics, after which
        // it must not be invocable again.
        self.state = State::Unavailable;

        // Switch to the generator function, and retrieve the yielded value.
        let val = unsafe {
          let mut data_in = (&mut self.context as *mut Context<Stack>, input);
          let data_out =
            ptr::read(Context::swap(&mut self.context, &self.context,
                                    &mut data_in as *mut (*mut Context<Stack>, Input)  as usize)
                      as *mut Option<Output>);
          mem::forget(data_in);
          data_out
        };

        // Unless the generator function has returned, it can be switched to again, so
        // set the state to Runnable.
        if val.is_some() { self.state = State::Runnable }

        val
      }
      State::Unavailable => None
    }
  }

  /// Returns the state of the generator.
  #[inline]
  pub fn state(&self) -> State { self.state }

  /// Extracts the stack from a generator when the generator function has returned.
  /// If the generator function has not returned
  /// (i.e. `self.state() == State::Runnable`), panics.
  pub fn unwrap(self) -> Stack {
    match self.state {
      State::Runnable    => panic!("Argh! Bastard! Don't touch that!"),
      State::Unavailable => unsafe { self.context.unwrap() }
    }
  }
}

/// Yielder is an interface provided to every generator through which it
/// returns a value.
#[derive(Debug)]
pub struct Yielder<Input: Send, Output: Send, Stack: stack::Stack> {
  context: *mut Context<Stack>,
  phantom: (PhantomData<*const Input>, PhantomData<*const Output>)
}

impl<Input, Output, Stack> Yielder<Input, Output, Stack>
    where Input: Send, Output: Send, Stack: stack::Stack {
  fn new(context: *mut Context<Stack>) -> Yielder<Input, Output, Stack> {
    Yielder {
      context: context,
      phantom: (PhantomData, PhantomData)
    }
  }

  #[inline(always)]
  fn suspend_bare(&mut self, mut val: Option<Output>) -> Input {
    unsafe {
      let data = Context::swap(self.context, self.context,
                               &mut val as *mut Option<Output> as usize);
      mem::forget(val);
      let (new_context, input) = ptr::read(data as *mut (*mut Context<Stack>, Input));
      // The generator can be moved (and with it, the context).
      // This changes the address of the context.
      // Thus, we update it after each swap.
      self.context = new_context;
      // However, between this point and the next time we enter suspend_bare
      // the generator cannot be moved, as a &mut Generator is necessary
      // to resume the generator function.
      input
    }
  }

  /// Suspends the generator and returns `Some(item)` from the `resume()`
  /// invocation that resumed the generator.
  #[inline(always)]
  pub fn suspend(&mut self, item: Output) -> Input {
    self.suspend_bare(Some(item))
  }
}

impl<Output, Stack> Iterator for Generator<(), Output, Stack>
    where Output: Send, Stack: stack::Stack {
  type Item = Output;

  fn next(&mut self) -> Option<Self::Item> { self.resume(()) }
}
<commit_msg>Yielder: use Cell internally to allow suspending via multiple borrows.<commit_after>// This file is part of libfringe, a low-level green threading library.
// Copyright (c) whitequark <whitequark@whitequark.org>
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

//! Generators.
//!
//! Generators allow repeatedly suspending the execution of a function,
//! returning a value to the caller, and resuming the suspended function
//! afterwards.

use core::marker::PhantomData;
use core::{ptr, mem};
use core::cell::Cell;

use stack;
use context::Context;

#[derive(Debug, Clone, Copy)]
pub enum State {
  /// Generator can be resumed. This is the initial state.
  Runnable,
  /// Generator cannot be resumed. This is the state of the generator after
  /// the generator function has returned or panicked.
  Unavailable
}

/// Generator wraps a function and allows suspending its execution more than once, returning
/// a value each time.
///
/// The first time `resume(input0)` is called, the function is called as `f(yielder, input0)`.
/// It runs until it suspends its execution through `yielder.suspend(output0)`, after which
/// `resume(input0)` returns `output0`. The function can be resumed again using `resume(input1)`,
/// after which `yielder.suspend(output0)` returns `input1`, and so on. Once the function returns,
/// the `resume()` call will return `None`, and it will return `None` every time it is called
/// after that.
///
/// If the generator function panics, the panic is propagated through the `resume()` call as usual.
///
/// After the generator function returns or panics, it is safe to reclaim the generator stack
/// using `unwrap()`.
///
/// `state()` can be used to determine whether the generator function has returned;
/// the state is `State::Runnable` after creation and suspension, and `State::Unavailable`
/// once the generator function returns or panics.
///
/// When the input type is `()`, a generator implements the Iterator trait.
///
/// # Example
///
/// ```
/// use fringe::{OsStack, Generator};
///
/// let stack = OsStack::new(0).unwrap();
/// let mut add_one = Generator::new(stack, move |yielder, mut input| {
///   loop {
///     if input == 0 { break }
///     input = yielder.suspend(input + 1)
///   }
/// });
/// println!("{:?}", add_one.resume(2)); // prints Some(3)
/// println!("{:?}", add_one.resume(3)); // prints Some(4)
/// println!("{:?}", add_one.resume(0)); // prints None
/// ```
///
/// # Iterator example
///
/// ```
/// use fringe::{OsStack, Generator};
///
/// let stack = OsStack::new(0).unwrap();
/// let mut nat = Generator::new(stack, move |yielder, ()| {
///   for i in 1.. { yielder.suspend(i) }
/// });
/// println!("{:?}", nat.next()); // prints Some(0)
/// println!("{:?}", nat.next()); // prints Some(1)
/// println!("{:?}", nat.next()); // prints Some(2)
/// ```
#[derive(Debug)]
pub struct Generator<Input: Send, Output: Send, Stack: stack::Stack> {
  state:   State,
  context: Context<Stack>,
  phantom: (PhantomData<*const Input>, PhantomData<*const Output>)
}

impl<Input, Output, Stack> Generator<Input, Output, Stack>
    where Input: Send, Output: Send, Stack: stack::Stack {
  /// Creates a new generator.
  ///
  /// See also the [contract](../trait.GuardedStack.html) that needs to be fulfilled by `stack`.
  pub fn new<F>(stack: Stack, f: F) -> Generator<Input, Output, Stack>
      where Stack: stack::GuardedStack,
            F: FnOnce(&mut Yielder<Input, Output, Stack>, Input) + Send {
    unsafe { Generator::unsafe_new(stack, f) }
  }

  /// Same as `new`, but does not require `stack` to have a guard page.
  ///
  /// This function is unsafe because the generator function can easily violate
  /// memory safety by overflowing the stack. It is useful in environments where
  /// guarded stacks do not exist, e.g. in absence of an MMU.
  ///
  /// See also the [contract](../trait.Stack.html) that needs to be fulfilled by `stack`.
  pub unsafe fn unsafe_new<F>(stack: Stack, f: F) -> Generator<Input, Output, Stack>
      where F: FnOnce(&mut Yielder<Input, Output, Stack>, Input) + Send {
    unsafe extern "C" fn generator_wrapper<Input, Output, Stack, F>(env: usize) -> !
        where Input: Send, Output: Send, Stack: stack::Stack,
              F: FnOnce(&mut Yielder<Input, Output, Stack>, Input) {
      // Retrieve our environment from the callee and return control to it.
      let (mut yielder, f) = ptr::read(env as *mut (Yielder<Input, Output, Stack>, F));
      let data = Context::swap(yielder.context.get(), yielder.context.get(), 0);
      // See the second half of Yielder::suspend_bare.
      let (new_context, input) = ptr::read(data as *mut (*mut Context<Stack>, Input));
      yielder.context.set(new_context as *mut Context<Stack>);
      // Run the body of the generator.
      f(&mut yielder, input);
      // Past this point, the generator has dropped everything it has held.
      loop { yielder.suspend_bare(None); }
    }

    let mut generator = Generator {
      state:   State::Runnable,
      context: Context::new(stack, generator_wrapper::<Input, Output, Stack, F>),
      phantom: (PhantomData, PhantomData)
    };

    // Transfer environment to the callee.
    let mut env = (Yielder::new(&mut generator.context), f);
    Context::swap(&mut generator.context, &generator.context,
                  &mut env as *mut (Yielder<Input, Output, Stack>, F) as usize);
    mem::forget(env);

    generator
  }

  /// Resumes the generator and return the next value it yields.
  /// If the generator function has returned, returns `None`.
  #[inline]
  pub fn resume(&mut self, input: Input) -> Option<Output> {
    match self.state {
      State::Runnable => {
        // Set the state to Unavailable. Since we have exclusive access to the generator,
        // the only case where this matters is the generator function panics, after which
        // it must not be invocable again.
        self.state = State::Unavailable;

        // Switch to the generator function, and retrieve the yielded value.
        let val = unsafe {
          let mut data_in = (&mut self.context as *mut Context<Stack>, input);
          let data_out =
            ptr::read(Context::swap(&mut self.context, &self.context,
                                    &mut data_in as *mut (*mut Context<Stack>, Input)  as usize)
                      as *mut Option<Output>);
          mem::forget(data_in);
          data_out
        };

        // Unless the generator function has returned, it can be switched to again, so
        // set the state to Runnable.
        if val.is_some() { self.state = State::Runnable }

        val
      }
      State::Unavailable => None
    }
  }

  /// Returns the state of the generator.
  #[inline]
  pub fn state(&self) -> State { self.state }

  /// Extracts the stack from a generator when the generator function has returned.
  /// If the generator function has not returned
  /// (i.e. `self.state() == State::Runnable`), panics.
  pub fn unwrap(self) -> Stack {
    match self.state {
      State::Runnable    => panic!("Argh! Bastard! Don't touch that!"),
      State::Unavailable => unsafe { self.context.unwrap() }
    }
  }
}

/// Yielder is an interface provided to every generator through which it
/// returns a value.
#[derive(Debug)]
pub struct Yielder<Input: Send, Output: Send, Stack: stack::Stack> {
  context: Cell<*mut Context<Stack>>,
  phantom: (PhantomData<*const Input>, PhantomData<*const Output>)
}

impl<Input, Output, Stack> Yielder<Input, Output, Stack>
    where Input: Send, Output: Send, Stack: stack::Stack {
  fn new(context: *mut Context<Stack>) -> Yielder<Input, Output, Stack> {
    Yielder {
      context: Cell::new(context),
      phantom: (PhantomData, PhantomData)
    }
  }

  #[inline(always)]
  fn suspend_bare(&self, mut val: Option<Output>) -> Input {
    unsafe {
      let data = Context::swap(self.context.get(), self.context.get(),
                               &mut val as *mut Option<Output> as usize);
      mem::forget(val);
      let (new_context, input) = ptr::read(data as *mut (*mut Context<Stack>, Input));
      // The generator can be moved (and with it, the context).
      // This changes the address of the context.
      // Thus, we update it after each swap.
      self.context.set(new_context);
      // However, between this point and the next time we enter suspend_bare
      // the generator cannot be moved, as a &mut Generator is necessary
      // to resume the generator function.
      input
    }
  }

  /// Suspends the generator and returns `Some(item)` from the `resume()`
  /// invocation that resumed the generator.
  #[inline(always)]
  pub fn suspend(&self, item: Output) -> Input {
    self.suspend_bare(Some(item))
  }
}

impl<Output, Stack> Iterator for Generator<(), Output, Stack>
    where Output: Send, Stack: stack::Stack {
  type Item = Output;

  fn next(&mut self) -> Option<Self::Item> { self.resume(()) }
}
<|endoftext|>


<commit_before>use std::collections::BTreeMap;
use std::path::{PathBuf, Path};
use config::Config;

#[derive(Debug, Clone, Default, Hash, PartialEq, Eq, PartialOrd)]
pub struct Trace {
    pub line: u64,
    pub address: Option<u64>,
    pub hits: u64,
}

/// Stores all the program traces mapped to files and provides an interface to
/// add, query and change traces.
pub struct TraceMap<'a> {
    traces: BTreeMap<PathBuf, Vec<Trace>>,
    config: &'a Config
}

impl<'a> TraceMap<'a> {

    pub fn new(config: &'a Config) -> TraceMap<'a> {
        TraceMap {
            config: config,
            traces: BTreeMap::new(),
        }
    }

    pub fn add_trace(&mut self, file: &Path, trace: Trace) {
        if self.traces.contains_key(file) {
            if let Some(trace_vec) = self.traces.get_mut(file) {
                trace_vec.push(trace);
            }
        } else {
            self.traces.insert(file.to_path_buf(), vec![trace]);
        }
    }

    pub fn get_trace(&self, address: u64) -> Option<&Trace> {
        self.all_traces()
            .iter()
            .find(|x| x.address == Some(address))
            .map(|x| *x)
    }

    pub fn get_trace_mut(&mut self, address: u64) -> Option<&mut Trace> {
        None
    }

    pub fn get_child_traces(&self, root: &Path) -> Vec<&Trace> {
        self.traces.iter()
                   .filter(|&(ref k, _)| k.starts_with(root))
                   .flat_map(|(_, ref v)| v.iter())
                   .collect()
    }

    fn all_traces(&self) -> Vec<&Trace> {
        self.traces.values().flat_map(|ref x| x.iter()).collect()
    }

}
<commit_msg>Refined coverage data structures<commit_after>use std::collections::BTreeMap;
use std::path::{PathBuf, Path};
use config::Config;

/// Used to track the state of logical conditions
#[derive(Debug, Clone, Default, Hash, PartialEq, Eq, PartialOrd)]
pub struct LogicState {
    /// Whether the condition has been observed as true
    pub beenTrue: bool,
    /// Whether the condition has been observed as false
    pub beenFalse: bool,
}

/// Shows what type of coverage data is being collected by a given trace
#[derive(Debug, Clone, Default, Hash, PartialEq, Eq, PartialOrd)]
pub enum CoverageStat {
    /// Line coverage data (whether line has been hit)
    Line(u64),
    /// Branch coverage data (whether branch has been true and false
    Branch(LogicState),
    /// Condition coverage data (each boolean subcondition true and false)
    Condition(Vec<LogicState>),
}

#[derive(Debug, Clone, Default, Hash, PartialEq, Eq, PartialOrd)]
pub struct Trace {
    /// Line the trace is on in the file
    pub line: u64,
    /// Optional address showing location in the test artefact
    pub address: Option<u64>,
    /// Length of the instruction (useful to get entire condition/branch)
    pub length: usize,
    /// Coverage stats
    pub stats: CoverageStat,
}

/// Stores all the program traces mapped to files and provides an interface to
/// add, query and change traces.
pub struct TraceMap<'a> {
    /// Traces in the program mapped to the given file
    traces: BTreeMap<PathBuf, Vec<Trace>>,
    /// Tarpaulin config - is this needed?
    config: &'a Config
}

impl<'a> TraceMap<'a> {
    /// Create a new TraceMap
    pub fn new(config: &'a Config) -> TraceMap<'a> {
        TraceMap {
            config: config,
            traces: BTreeMap::new(),
        }
    }

    /// Add a trace to the tracemap for the given file
    pub fn add_trace(&mut self, file: &Path, trace: Trace) {
        if self.traces.contains_key(file) {
            if let Some(trace_vec) = self.traces.get_mut(file) {
                trace_vec.push(trace);
            }
        } else {
            self.traces.insert(file.to_path_buf(), vec![trace]);
        }
    }

    /// Gets an immutable reference to a trace from an address. Returns None if
    /// there is no trace at that address
    pub fn get_trace(&self, address: u64) -> Option<&Trace> {
        self.all_traces()
            .iter()
            .find(|x| x.address == Some(address))
            .map(|x| *x)
    }

    /// Gets a mutable reference to a trace at a given address
    /// Returns None if there is no trace at that address
    pub fn get_trace_mut(&mut self, address: u64) -> Option<&mut Trace> {
        None
    }

    /// Gets all traces below a certain path
    pub fn get_child_traces(&self, root: &Path) -> Vec<&Trace> {
        self.traces.iter()
                   .filter(|&(ref k, _)| k.starts_with(root))
                   .flat_map(|(_, ref v)| v.iter())
                   .collect()
    }

    /// Gets all traces
    fn all_traces(&self) -> Vec<&Trace> {
        self.traces.values().flat_map(|ref x| x.iter()).collect()
    }

}
<|endoftext|>


<commit_before>use bytes::Bytes;
use futures::{Poll, Stream};
use futures::sync::mpsc;
use tokio_proto;
use std::borrow::Cow;

use http::Chunk;

pub type TokioBody = tokio_proto::streaming::Body<Chunk, ::Error>;

/// A `Stream` for `Chunk`s used in requests and responses.
#[must_use = "streams do nothing unless polled"]
#[derive(Debug)]
pub struct Body(TokioBody);

impl Body {
    /// Return an empty body stream
    #[inline]
    pub fn empty() -> Body {
        Body(TokioBody::empty())
    }

    /// Return a body stream with an associated sender half
    #[inline]
    pub fn pair() -> (mpsc::Sender<Result<Chunk, ::Error>>, Body) {
        let (tx, rx) = TokioBody::pair();
        let rx = Body(rx);
        (tx, rx)
    }
}

impl Default for Body {
    #[inline]
    fn default() -> Body {
        Body::empty()
    }
}

impl Stream for Body {
    type Item = Chunk;
    type Error = ::Error;

    #[inline]
    fn poll(&mut self) -> Poll<Option<Chunk>, ::Error> {
        self.0.poll()
    }
}

impl From<Body> for tokio_proto::streaming::Body<Chunk, ::Error> {
    #[inline]
    fn from(b: Body) -> tokio_proto::streaming::Body<Chunk, ::Error> {
        b.0
    }
}

impl From<tokio_proto::streaming::Body<Chunk, ::Error>> for Body {
    #[inline]
    fn from(tokio_body: tokio_proto::streaming::Body<Chunk, ::Error>) -> Body {
        Body(tokio_body)
    }
}

impl From<mpsc::Receiver<Result<Chunk, ::Error>>> for Body {
    #[inline]
    fn from(src: mpsc::Receiver<Result<Chunk, ::Error>>) -> Body {
        Body(src.into())
    }
}

impl From<Chunk> for Body {
    #[inline]
    fn from (chunk: Chunk) -> Body {
        Body(TokioBody::from(chunk))
    }
}

impl From<Bytes> for Body {
    #[inline]
    fn from (bytes: Bytes) -> Body {
        Body(TokioBody::from(Chunk::from(bytes)))
    }
}

impl From<Vec<u8>> for Body {
    #[inline]
    fn from (vec: Vec<u8>) -> Body {
        Body(TokioBody::from(Chunk::from(vec)))
    }
}

impl From<&'static [u8]> for Body {
    #[inline]
    fn from (slice: &'static [u8]) -> Body {
        Body(TokioBody::from(Chunk::from(slice)))
    }
}

impl From<Cow<'static, [u8]>> for Body {
    #[inline]
    fn from (cow: Cow<'static, [u8]>) -> Body {
		if let Cow::Borrowed(value) = cow {
			Body::from(value)
		} else {
			Body::from(cow.to_owned())
		}
    }
}

impl From<String> for Body {
    #[inline]
    fn from (s: String) -> Body {
        Body(TokioBody::from(Chunk::from(s.into_bytes())))
    }
}

impl From<&'static str> for Body {
    #[inline]
    fn from (slice: &'static str) -> Body {
        Body(TokioBody::from(Chunk::from(slice.as_bytes())))
    }
}

impl From<Cow<'static, str>> for Body {
    #[inline]
    fn from (cow: Cow<'static, str>) -> Body {
		if let Cow::Borrowed(value) = cow {
			Body::from(value)
		} else {
			Body::from(cow.to_owned())
		}
    }
}

impl From<Option<Body>> for Body {
    #[inline]
    fn from (body: Option<Body>) -> Body {
        body.unwrap_or_default()
    }
}

fn _assert_send_sync() {
    fn _assert_send<T: Send>() {}
    fn _assert_sync<T: Sync>() {}

    _assert_send::<Body>();
    _assert_send::<Chunk>();
    _assert_sync::<Chunk>();
}

#[test]
fn test_body_stream_concat() {
    use futures::{Sink, Stream, Future};
    let (tx, body) = Body::pair();

    ::std::thread::spawn(move || {
        let tx = tx.send(Ok("hello ".into())).wait().unwrap();
        tx.send(Ok("world".into())).wait().unwrap();
    });

    let total = body.concat2().wait().unwrap();
    assert_eq!(total.as_ref(), b"hello world");

}
<commit_msg>style(body): change tabs to spaces<commit_after>use bytes::Bytes;
use futures::{Poll, Stream};
use futures::sync::mpsc;
use tokio_proto;
use std::borrow::Cow;

use http::Chunk;

pub type TokioBody = tokio_proto::streaming::Body<Chunk, ::Error>;

/// A `Stream` for `Chunk`s used in requests and responses.
#[must_use = "streams do nothing unless polled"]
#[derive(Debug)]
pub struct Body(TokioBody);

impl Body {
    /// Return an empty body stream
    #[inline]
    pub fn empty() -> Body {
        Body(TokioBody::empty())
    }

    /// Return a body stream with an associated sender half
    #[inline]
    pub fn pair() -> (mpsc::Sender<Result<Chunk, ::Error>>, Body) {
        let (tx, rx) = TokioBody::pair();
        let rx = Body(rx);
        (tx, rx)
    }
}

impl Default for Body {
    #[inline]
    fn default() -> Body {
        Body::empty()
    }
}

impl Stream for Body {
    type Item = Chunk;
    type Error = ::Error;

    #[inline]
    fn poll(&mut self) -> Poll<Option<Chunk>, ::Error> {
        self.0.poll()
    }
}

impl From<Body> for tokio_proto::streaming::Body<Chunk, ::Error> {
    #[inline]
    fn from(b: Body) -> tokio_proto::streaming::Body<Chunk, ::Error> {
        b.0
    }
}

impl From<tokio_proto::streaming::Body<Chunk, ::Error>> for Body {
    #[inline]
    fn from(tokio_body: tokio_proto::streaming::Body<Chunk, ::Error>) -> Body {
        Body(tokio_body)
    }
}

impl From<mpsc::Receiver<Result<Chunk, ::Error>>> for Body {
    #[inline]
    fn from(src: mpsc::Receiver<Result<Chunk, ::Error>>) -> Body {
        Body(src.into())
    }
}

impl From<Chunk> for Body {
    #[inline]
    fn from (chunk: Chunk) -> Body {
        Body(TokioBody::from(chunk))
    }
}

impl From<Bytes> for Body {
    #[inline]
    fn from (bytes: Bytes) -> Body {
        Body(TokioBody::from(Chunk::from(bytes)))
    }
}

impl From<Vec<u8>> for Body {
    #[inline]
    fn from (vec: Vec<u8>) -> Body {
        Body(TokioBody::from(Chunk::from(vec)))
    }
}

impl From<&'static [u8]> for Body {
    #[inline]
    fn from (slice: &'static [u8]) -> Body {
        Body(TokioBody::from(Chunk::from(slice)))
    }
}

impl From<Cow<'static, [u8]>> for Body {
    #[inline]
    fn from (cow: Cow<'static, [u8]>) -> Body {
        if let Cow::Borrowed(value) = cow {
            Body::from(value)
        } else {
            Body::from(cow.to_owned())
        }
    }
}

impl From<String> for Body {
    #[inline]
    fn from (s: String) -> Body {
        Body(TokioBody::from(Chunk::from(s.into_bytes())))
    }
}

impl From<&'static str> for Body {
    #[inline]
    fn from(slice: &'static str) -> Body {
        Body(TokioBody::from(Chunk::from(slice.as_bytes())))
    }
}

impl From<Cow<'static, str>> for Body {
    #[inline]
    fn from(cow: Cow<'static, str>) -> Body {
        if let Cow::Borrowed(value) = cow {
            Body::from(value)
        } else {
            Body::from(cow.to_owned())
        }
    }
}

impl From<Option<Body>> for Body {
    #[inline]
    fn from (body: Option<Body>) -> Body {
        body.unwrap_or_default()
    }
}

fn _assert_send_sync() {
    fn _assert_send<T: Send>() {}
    fn _assert_sync<T: Sync>() {}

    _assert_send::<Body>();
    _assert_send::<Chunk>();
    _assert_sync::<Chunk>();
}

#[test]
fn test_body_stream_concat() {
    use futures::{Sink, Stream, Future};
    let (tx, body) = Body::pair();

    ::std::thread::spawn(move || {
        let tx = tx.send(Ok("hello ".into())).wait().unwrap();
        tx.send(Ok("world".into())).wait().unwrap();
    });

    let total = body.concat2().wait().unwrap();
    assert_eq!(total.as_ref(), b"hello world");

}
<|endoftext|>


<commit_before><commit_msg>hello world<commit_after><|endoftext|>


<commit_before>//! Utility functions to write nifti images.

use std::fs::File;
use std::io::{BufWriter, Write};
use std::mem;
use std::ops::{Div, Sub};
use std::slice::from_raw_parts;

use byteorder::{LittleEndian, WriteBytesExt};
use flate2::Compression;
use flate2::write::GzEncoder;
use ndarray::{ArrayBase, Axis, Data, Dimension, RemoveAxis, ScalarOperand};
use num_traits::FromPrimitive;

use {NiftiHeader, volume::element::DataElement, util::is_gz_file};

type B = LittleEndian;

/// Write a nifti file (.nii or .nii.gz).
///
/// If a `reference` is given, it will be used to fill most of the header's fields. The voxels
/// intensity will be subtracted by `scl_slope` and divided by `scl_inter`. If `reference` is not
/// given, a default `NiftiHeader` will be built and written.
///
/// In all cases, the `dim`, `datatype` and `bitpix` fields will depend only on `data`, not on the
/// header. In other words, the `datatype` defined in `reference` will be ignored.
pub fn write_nifti<A, S, D>(
    path: &str,
    data: &ArrayBase<S, D>,
    reference: Option<&NiftiHeader>)
    where S: Data<Elem=A>,
          A: Copy + DataElement + Div<Output = A> + FromPrimitive + ScalarOperand + Sub<Output = A>,
          D: Dimension + RemoveAxis
{
    let mut dim = [1; 8];
    dim[0] = data.ndim() as u16;
    for (i, s) in data.shape().iter().enumerate() {
        dim[i + 1] = *s as u16;
    }

    // If no reference header is given, use the default.
    let reference = match reference {
        Some(r) => r.clone(),
        None => NiftiHeader::default()
    };
    let header = NiftiHeader {
        dim,
        datatype: A::DATA_TYPE as i16,
        bitpix: (mem::size_of::<A>() * 8) as i16,
        // All other fields are copied from reference header
        ..reference
    };

    let f = File::create(path).expect("Can't create new nifti file");
    let mut writer = BufWriter::new(f);
    if is_gz_file(&path) {
        let mut buffer = Vec::new();
        write(&mut buffer, header, data);

        let mut e = GzEncoder::new(Vec::new(), Compression::default());
        e.write_all(&buffer).unwrap();

        let buffer = e.finish().unwrap();
        writer.write_all(&buffer).unwrap();
    } else {
        write(&mut writer, header, data);
    }
}

fn write<A, S, D, W>(
    writer: &mut W,
    header: NiftiHeader,
    data: &ArrayBase<S, D>,)
    where S: Data<Elem=A>,
          A: Copy + DataElement + Div<Output = A> + FromPrimitive + ScalarOperand + Sub<Output = A>,
          D: Dimension + RemoveAxis,
          W: WriteBytesExt
{
    writer.write_i32::<B>(header.sizeof_hdr).unwrap();
    writer.write_all(&header.data_type).unwrap();
    writer.write_all(&header.db_name).unwrap();
    writer.write_i32::<B>(header.extents).unwrap();
    writer.write_i16::<B>(header.session_error).unwrap();
    writer.write_u8(header.regular).unwrap();
    writer.write_u8(header.dim_info).unwrap();
    for s in &header.dim {
        writer.write_u16::<B>(*s).unwrap();
    }
    writer.write_f32::<B>(header.intent_p1).unwrap();
    writer.write_f32::<B>(header.intent_p2).unwrap();
    writer.write_f32::<B>(header.intent_p3).unwrap();
    writer.write_i16::<B>(header.intent_code).unwrap();
    writer.write_i16::<B>(header.datatype).unwrap();
    writer.write_i16::<B>(header.bitpix).unwrap();
    writer.write_i16::<B>(header.slice_start).unwrap();
    for f in &header.pixdim {
        writer.write_f32::<B>(*f).unwrap();
    }
    writer.write_f32::<B>(header.vox_offset).unwrap();
    writer.write_f32::<B>(header.scl_slope).unwrap();
    writer.write_f32::<B>(header.scl_inter).unwrap();
    writer.write_i16::<B>(header.slice_end).unwrap();
    writer.write_u8(header.slice_code).unwrap();
    writer.write_u8(header.xyzt_units).unwrap();
    writer.write_f32::<B>(header.cal_max).unwrap();
    writer.write_f32::<B>(header.cal_min).unwrap();
    writer.write_f32::<B>(header.slice_duration).unwrap();
    writer.write_f32::<B>(header.toffset).unwrap();
    writer.write_i32::<B>(header.glmax).unwrap();
    writer.write_i32::<B>(header.glmin).unwrap();

    writer.write_all(&header.descrip).unwrap();
    writer.write_all(&header.aux_file).unwrap();
    writer.write_i16::<B>(header.qform_code).unwrap();
    writer.write_i16::<B>(header.sform_code).unwrap();
    for f in &[header.quatern_b, header.quatern_c, header.quatern_d,
               header.quatern_x, header.quatern_y, header.quatern_z] {
        writer.write_f32::<B>(*f).unwrap();
    }
    for f in header.srow_x.iter().chain(&header.srow_y).chain(&header.srow_z) {
        writer.write_f32::<B>(*f).unwrap();
    }
    writer.write_all(&header.intent_name).unwrap();
    writer.write_all(&header.magic).unwrap();

    // Empty 4 bytes after the header
    writer.write_u32::<B>(0).unwrap();

    // We finally write the data.
    // Like NiBabel, we iterate by "slice" to improve speed and use less memory

    // Need the transpose for fortran used in nifti file format.
    let data = data.t();

    let arr_len = data.subview(Axis(0), 0).len();
    let resolution: usize = header.dim[1..(header.dim[0] + 1) as usize]
        .iter().map(|d| *d as usize).product();
    let nb_bytes = resolution * header.bitpix as usize / 8;
    let dim0_size = data.shape()[0];

    // `1.0x + 0.0` would give the same results, but we avoid a lot of divisions
    let slope = if header.scl_slope == 0.0 { 1.0 } else { header.scl_slope };
    if slope != 1.0 || header.scl_inter != 0.0 {
        let slope = A::from_f32(slope).unwrap();
        let inter = A::from_f32(header.scl_inter).unwrap();
        for arr_data in data.axis_iter(Axis(0)) {
            let arr_data = arr_data.sub(inter).div(slope);
            let arr_data = arr_data.into_shape(arr_len).unwrap();
            let arr_data = arr_data.as_slice().unwrap();

            let buffer = to_bytes(&arr_data, nb_bytes / dim0_size);
            writer.write_all(&buffer).unwrap();
        }
    } else {
        for arr_data in data.axis_iter(Axis(0)) {
            // We need to own the data because of the into_shape() for `C` ordering.
            let arr_data = arr_data.to_owned();
            let arr_data = arr_data.into_shape(arr_len).unwrap();
            let arr_data = arr_data.as_slice().unwrap();

            let buffer = to_bytes(&arr_data, nb_bytes / dim0_size);
            writer.write_all(&buffer).unwrap();
        }
    }
}

fn to_bytes<T>(data: &[T], size: usize) -> &[u8] {
    unsafe { from_raw_parts::<u8>(data.as_ptr() as *const u8, size) }
}

#[cfg(test)]
pub mod tests {
    extern crate tempfile;

    use super::*;

    use ndarray::{Array, Array2, Ix2, IxDyn, ShapeBuilder};
    use self::tempfile::tempdir;

    use {InMemNiftiObject, IntoNdArray, header::MAGIC_CODE_NIP1, object::NiftiObject};

    fn get_random_path(ext: &str) -> String {
        let dir = tempdir().unwrap();
        let path = if ext == "" {
            dir.into_path()
        } else {
            dir.into_path().join(ext)
        };
        path.to_str().unwrap().to_string()
    }

    pub fn generate_nifti_header(
        dim: [u16; 8],
        scl_slope: f32,
        scl_inter: f32,
        datatype: i16
    ) -> NiftiHeader {
        let bitpix = (mem::size_of::<f32>() * 8) as i16;
        let magic = *MAGIC_CODE_NIP1;
        NiftiHeader {
            dim, datatype, bitpix, magic, scl_slope, scl_inter, ..NiftiHeader::default()
        }
    }

    fn read_2d_image(path: &str) -> Array2<f32> {
        let nifti_object = InMemNiftiObject::from_file(path).expect("Nifti file is unreadable.");
        let volume = nifti_object.into_volume();
        let dyn_data = volume.into_ndarray().unwrap();
        dyn_data.into_dimensionality::<Ix2>().unwrap()
    }

    fn test_write_read(arr: &Array<f32, IxDyn>, path: &str) {
        let path = get_random_path(path);
        let mut dim = [1; 8];
        dim[0] = arr.ndim() as u16;
        for (i, s) in arr.shape().iter().enumerate() {
            dim[i + 1] = *s as u16;
        }
        let header = generate_nifti_header(dim, 1.0, 0.0, 16);
        write_nifti(&path, &arr, Some(&header));

        let read_nifti = read_2d_image(&path);
        assert!(read_nifti.all_close(&arr, 1e-10) == true);
    }

    fn f_order_array() -> Array<f32, IxDyn> {
        let dim = vec![4, 4];
        let vec = (0..16).map(|x| x as f32).collect();
        Array::from_shape_vec(IxDyn(&dim).f(), vec).unwrap()
    }

    fn c_order_array() -> Array<f32, IxDyn> {
        let dim = vec![4, 4];
        let vec = (0..16).map(|x| x as f32).collect();
        Array::from_shape_vec(IxDyn(&dim), vec).unwrap()
    }

    #[test]
    fn test_fortran_writing() {
        // Test .nii
        let arr = f_order_array();
        test_write_read(&arr, "test.nii");
        let mut arr = f_order_array();
        arr.invert_axis(Axis(1));
        test_write_read(&arr, "test_non_contiguous.nii");

        // Test .nii.gz
        let arr = f_order_array();
        test_write_read(&arr, "test.nii.gz");
        let mut arr = f_order_array();
        arr.invert_axis(Axis(1));
        test_write_read(&arr, "test_non_contiguous.nii.gz");
    }

    #[test]
    fn test_c_writing() {
        // Test .nii
        let arr = c_order_array();
        test_write_read(&arr, "test.nii");

        let mut arr = c_order_array();
        arr.invert_axis(Axis(1));
        test_write_read(&arr, "test_non_contiguous.nii");

        // Test .nii.gz
        let arr = c_order_array();
        test_write_read(&arr, "test.nii.gz");

        let mut arr = c_order_array();
        arr.invert_axis(Axis(1));
        test_write_read(&arr, "test_non_contiguous.nii.gz");
    }

    #[test]
    fn test_header_slope_inter() {
        use std::ops::{Add, Mul};

        let arr = f_order_array();
        let slope = 2.2;
        let inter = 101.1;

        let path = get_random_path("test_slope_inter.nii");
        let mut dim = [1; 8];
        dim[0] = arr.ndim() as u16;
        for (i, s) in arr.shape().iter().enumerate() {
            dim[i + 1] = *s as u16;
        }
        let header = generate_nifti_header(dim, slope, inter, 16);
        let transformed_data = arr.mul(slope).add(inter);
        write_nifti(&path, &transformed_data, Some(&header));

        let read_nifti = read_2d_image(&path);
        assert!(read_nifti.all_close(&transformed_data, 1e-10) == true);
    }
}
<commit_msg>Move header writing to a separate function<commit_after>//! Utility functions to write nifti images.

use std::fs::File;
use std::io::{BufWriter, Write};
use std::mem;
use std::ops::{Div, Sub};
use std::slice::from_raw_parts;

use byteorder::{LittleEndian, WriteBytesExt};
use flate2::Compression;
use flate2::write::GzEncoder;
use ndarray::{ArrayBase, Axis, Data, Dimension, RemoveAxis, ScalarOperand};
use num_traits::FromPrimitive;

use {NiftiHeader, volume::element::DataElement, util::is_gz_file};

type B = LittleEndian;

/// Write a nifti file (.nii or .nii.gz).
///
/// If a `reference` is given, it will be used to fill most of the header's fields. The voxels
/// intensity will be subtracted by `scl_slope` and divided by `scl_inter`. If `reference` is not
/// given, a default `NiftiHeader` will be built and written.
///
/// In all cases, the `dim`, `datatype` and `bitpix` fields will depend only on `data`, not on the
/// header. In other words, the `datatype` defined in `reference` will be ignored.
pub fn write_nifti<A, S, D>(
    path: &str,
    data: &ArrayBase<S, D>,
    reference: Option<&NiftiHeader>)
    where S: Data<Elem=A>,
          A: Copy + DataElement + Div<Output = A> + FromPrimitive + ScalarOperand + Sub<Output = A>,
          D: Dimension + RemoveAxis
{
    let mut dim = [1; 8];
    dim[0] = data.ndim() as u16;
    for (i, s) in data.shape().iter().enumerate() {
        dim[i + 1] = *s as u16;
    }

    // If no reference header is given, use the default.
    let reference = match reference {
        Some(r) => r.clone(),
        None => NiftiHeader::default()
    };
    let header = NiftiHeader {
        dim,
        datatype: A::DATA_TYPE as i16,
        bitpix: (mem::size_of::<A>() * 8) as i16,
        // All other fields are copied from reference header
        ..reference
    };

    let f = File::create(path).expect("Can't create new nifti file");
    let mut writer = BufWriter::new(f);
    if is_gz_file(&path) {
        let mut buffer = Vec::new();
        write_header(&mut buffer, &header);
        write_data(&mut buffer, header, data);

        let mut e = GzEncoder::new(Vec::new(), Compression::default());
        e.write_all(&buffer).unwrap();

        let buffer = e.finish().unwrap();
        writer.write_all(&buffer).unwrap();
    } else {
        write_header(&mut writer, &header);
        write_data(&mut writer, header, data);
    }
}

fn write_header<W>(
    writer: &mut W,
    header: &NiftiHeader)
    where W: WriteBytesExt
{
    writer.write_i32::<B>(header.sizeof_hdr).unwrap();
    writer.write_all(&header.data_type).unwrap();
    writer.write_all(&header.db_name).unwrap();
    writer.write_i32::<B>(header.extents).unwrap();
    writer.write_i16::<B>(header.session_error).unwrap();
    writer.write_u8(header.regular).unwrap();
    writer.write_u8(header.dim_info).unwrap();
    for s in &header.dim {
        writer.write_u16::<B>(*s).unwrap();
    }
    writer.write_f32::<B>(header.intent_p1).unwrap();
    writer.write_f32::<B>(header.intent_p2).unwrap();
    writer.write_f32::<B>(header.intent_p3).unwrap();
    writer.write_i16::<B>(header.intent_code).unwrap();
    writer.write_i16::<B>(header.datatype).unwrap();
    writer.write_i16::<B>(header.bitpix).unwrap();
    writer.write_i16::<B>(header.slice_start).unwrap();
    for f in &header.pixdim {
        writer.write_f32::<B>(*f).unwrap();
    }
    writer.write_f32::<B>(header.vox_offset).unwrap();
    writer.write_f32::<B>(header.scl_slope).unwrap();
    writer.write_f32::<B>(header.scl_inter).unwrap();
    writer.write_i16::<B>(header.slice_end).unwrap();
    writer.write_u8(header.slice_code).unwrap();
    writer.write_u8(header.xyzt_units).unwrap();
    writer.write_f32::<B>(header.cal_max).unwrap();
    writer.write_f32::<B>(header.cal_min).unwrap();
    writer.write_f32::<B>(header.slice_duration).unwrap();
    writer.write_f32::<B>(header.toffset).unwrap();
    writer.write_i32::<B>(header.glmax).unwrap();
    writer.write_i32::<B>(header.glmin).unwrap();

    writer.write_all(&header.descrip).unwrap();
    writer.write_all(&header.aux_file).unwrap();
    writer.write_i16::<B>(header.qform_code).unwrap();
    writer.write_i16::<B>(header.sform_code).unwrap();
    for f in &[header.quatern_b, header.quatern_c, header.quatern_d,
               header.quatern_x, header.quatern_y, header.quatern_z] {
        writer.write_f32::<B>(*f).unwrap();
    }
    for f in header.srow_x.iter().chain(&header.srow_y).chain(&header.srow_z) {
        writer.write_f32::<B>(*f).unwrap();
    }
    writer.write_all(&header.intent_name).unwrap();
    writer.write_all(&header.magic).unwrap();

    // Empty 4 bytes after the header
    writer.write_u32::<B>(0).unwrap();
}

fn write_data<A, S, D, W>(
    writer: &mut W,
    header: NiftiHeader,
    data: &ArrayBase<S, D>)
    where S: Data<Elem=A>,
          A: Copy + DataElement + Div<Output = A> + FromPrimitive + ScalarOperand + Sub<Output = A>,
          D: Dimension + RemoveAxis,
          W: WriteBytesExt
{
    // We finally write the data.
    // Like NiBabel, we iterate by "slice" to improve speed and use less memory

    // Need the transpose for fortran used in nifti file format.
    let data = data.t();

    let arr_len = data.subview(Axis(0), 0).len();
    let resolution: usize = header.dim[1..(header.dim[0] + 1) as usize]
        .iter().map(|d| *d as usize).product();
    let nb_bytes = resolution * header.bitpix as usize / 8;
    let dim0_size = data.shape()[0];

    // `1.0x + 0.0` would give the same results, but we avoid a lot of divisions
    let slope = if header.scl_slope == 0.0 { 1.0 } else { header.scl_slope };
    if slope != 1.0 || header.scl_inter != 0.0 {
        let slope = A::from_f32(slope).unwrap();
        let inter = A::from_f32(header.scl_inter).unwrap();
        for arr_data in data.axis_iter(Axis(0)) {
            let arr_data = arr_data.sub(inter).div(slope);
            let arr_data = arr_data.into_shape(arr_len).unwrap();
            let arr_data = arr_data.as_slice().unwrap();

            let buffer = to_bytes(&arr_data, nb_bytes / dim0_size);
            writer.write_all(&buffer).unwrap();
        }
    } else {
        for arr_data in data.axis_iter(Axis(0)) {
            // We need to own the data because of the into_shape() for `C` ordering.
            let arr_data = arr_data.to_owned();
            let arr_data = arr_data.into_shape(arr_len).unwrap();
            let arr_data = arr_data.as_slice().unwrap();

            let buffer = to_bytes(&arr_data, nb_bytes / dim0_size);
            writer.write_all(&buffer).unwrap();
        }
    }
}

fn to_bytes<T>(data: &[T], size: usize) -> &[u8] {
    unsafe { from_raw_parts::<u8>(data.as_ptr() as *const u8, size) }
}

#[cfg(test)]
pub mod tests {
    extern crate tempfile;

    use super::*;

    use ndarray::{Array, Array2, Ix2, IxDyn, ShapeBuilder};
    use self::tempfile::tempdir;

    use {InMemNiftiObject, IntoNdArray, header::MAGIC_CODE_NIP1, object::NiftiObject};

    fn get_random_path(ext: &str) -> String {
        let dir = tempdir().unwrap();
        let path = if ext == "" {
            dir.into_path()
        } else {
            dir.into_path().join(ext)
        };
        path.to_str().unwrap().to_string()
    }

    pub fn generate_nifti_header(
        dim: [u16; 8],
        scl_slope: f32,
        scl_inter: f32,
        datatype: i16
    ) -> NiftiHeader {
        let bitpix = (mem::size_of::<f32>() * 8) as i16;
        let magic = *MAGIC_CODE_NIP1;
        NiftiHeader {
            dim, datatype, bitpix, magic, scl_slope, scl_inter, ..NiftiHeader::default()
        }
    }

    fn read_2d_image(path: &str) -> Array2<f32> {
        let nifti_object = InMemNiftiObject::from_file(path).expect("Nifti file is unreadable.");
        let volume = nifti_object.into_volume();
        let dyn_data = volume.into_ndarray().unwrap();
        dyn_data.into_dimensionality::<Ix2>().unwrap()
    }

    fn test_write_read(arr: &Array<f32, IxDyn>, path: &str) {
        let path = get_random_path(path);
        let mut dim = [1; 8];
        dim[0] = arr.ndim() as u16;
        for (i, s) in arr.shape().iter().enumerate() {
            dim[i + 1] = *s as u16;
        }
        let header = generate_nifti_header(dim, 1.0, 0.0, 16);
        write_nifti(&path, &arr, Some(&header));

        let read_nifti = read_2d_image(&path);
        assert!(read_nifti.all_close(&arr, 1e-10) == true);
    }

    fn f_order_array() -> Array<f32, IxDyn> {
        let dim = vec![4, 4];
        let vec = (0..16).map(|x| x as f32).collect();
        Array::from_shape_vec(IxDyn(&dim).f(), vec).unwrap()
    }

    fn c_order_array() -> Array<f32, IxDyn> {
        let dim = vec![4, 4];
        let vec = (0..16).map(|x| x as f32).collect();
        Array::from_shape_vec(IxDyn(&dim), vec).unwrap()
    }

    #[test]
    fn test_fortran_writing() {
        // Test .nii
        let arr = f_order_array();
        test_write_read(&arr, "test.nii");
        let mut arr = f_order_array();
        arr.invert_axis(Axis(1));
        test_write_read(&arr, "test_non_contiguous.nii");

        // Test .nii.gz
        let arr = f_order_array();
        test_write_read(&arr, "test.nii.gz");
        let mut arr = f_order_array();
        arr.invert_axis(Axis(1));
        test_write_read(&arr, "test_non_contiguous.nii.gz");
    }

    #[test]
    fn test_c_writing() {
        // Test .nii
        let arr = c_order_array();
        test_write_read(&arr, "test.nii");

        let mut arr = c_order_array();
        arr.invert_axis(Axis(1));
        test_write_read(&arr, "test_non_contiguous.nii");

        // Test .nii.gz
        let arr = c_order_array();
        test_write_read(&arr, "test.nii.gz");

        let mut arr = c_order_array();
        arr.invert_axis(Axis(1));
        test_write_read(&arr, "test_non_contiguous.nii.gz");
    }

    #[test]
    fn test_header_slope_inter() {
        use std::ops::{Add, Mul};

        let arr = f_order_array();
        let slope = 2.2;
        let inter = 101.1;

        let path = get_random_path("test_slope_inter.nii");
        let mut dim = [1; 8];
        dim[0] = arr.ndim() as u16;
        for (i, s) in arr.shape().iter().enumerate() {
            dim[i + 1] = *s as u16;
        }
        let header = generate_nifti_header(dim, slope, inter, 16);
        let transformed_data = arr.mul(slope).add(inter);
        write_nifti(&path, &transformed_data, Some(&header));

        let read_nifti = read_2d_image(&path);
        assert!(read_nifti.all_close(&transformed_data, 1e-10) == true);
    }
}
<|endoftext|>