<commit_before>#[crate_id = "bignum#0.1.0-pre"];

#[comment = "Bignum library for Rust"];
#[crate_type = "rlib"];

#[feature(macro_rules)];

extern crate gmp;

use gmp::{Mpz, RandState};
use std::fmt;
use std::num::ToStrRadix;
use std::rand::Rng;
use std::libc::c_ulong;

pub struct BigUint {
    data: Mpz
}

impl FromPrimitive for BigUint {
    fn from_u64(other: u64) -> Option<BigUint> {
        let mpz: Option<Mpz> = FromPrimitive::from_u64(other);
        match mpz {
            Some(mpz) => Some(BigUint{ data: mpz }),
            None      => None
        }
    }

    fn from_i64(other: i64) -> Option<BigUint> {
        if other < 0 {
            None
        } else {
            let mpz: Option<Mpz> = FromPrimitive::from_i64(other);
            match mpz {
                Some(mpz) => Some(BigUint{ data: mpz }),
                None      => None
            }
        }
    }
}

pub trait ToBigUint {
    fn to_biguint(&self) -> Option<BigUint>;
}

macro_rules! impl_to_biguint(
    ($T:ty, $from_ty:path) => {
        impl ToBigUint for $T {
            fn to_biguint(&self) -> Option<BigUint> {
                $from_ty(*self)
            }
        }
    }
)

impl_to_biguint!(int,  FromPrimitive::from_int)
impl_to_biguint!(i8,   FromPrimitive::from_i8)
impl_to_biguint!(i16,  FromPrimitive::from_i16)
impl_to_biguint!(i32,  FromPrimitive::from_i32)
impl_to_biguint!(i64,  FromPrimitive::from_i64)
impl_to_biguint!(uint, FromPrimitive::from_uint)
impl_to_biguint!(u8,   FromPrimitive::from_u8)
impl_to_biguint!(u16,  FromPrimitive::from_u16)
impl_to_biguint!(u32,  FromPrimitive::from_u32)
impl_to_biguint!(u64,  FromPrimitive::from_u64)

impl ToStrRadix for BigUint {
    fn to_str_radix(&self, radix: uint) -> ~str {
        self.data.to_str_radix(radix)
    }
}

impl fmt::Show for BigUint {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f.buf, "{}", self.to_str_radix(10))
    }
}

impl Add<BigUint, BigUint> for BigUint {
    fn add(&self, other: &BigUint) -> BigUint {
        BigUint{data: self.data.add(&other.data)}
    }
}

impl Sub<BigUint, BigUint> for BigUint {
    fn sub(&self, other: &BigUint) -> BigUint {
        BigUint{data: self.data.sub(&other.data)}
    }
}

impl Mul<BigUint, BigUint> for BigUint {
    fn mul(&self, other: &BigUint) -> BigUint {
        BigUint{data: self.data.mul(&other.data)}
    }
}

pub trait RandBigInt {
    /// Generate a random `BigUint` of the given bit size.
    fn gen_biguint(&mut self, bit_size: uint) -> BigUint;

    // /// Generate a random BigInt of the given bit size.
    // fn gen_bigint(&mut self, bit_size: uint) -> BigInt;

    /// Generate a random `BigUint` less than the given bound. Fails
    /// when the bound is zero.
    fn gen_biguint_below(&mut self, bound: &BigUint) -> BigUint;

    /// Generate a random `BigUint` within the given range. The lower
    /// bound is inclusive; the upper bound is exclusive. Fails when
    /// the upper bound is not greater than the lower bound.
    fn gen_biguint_range(&mut self, lbound: &BigUint, ubound: &BigUint) -> BigUint;

    // /// Generate a random `BigInt` within the given range. The lower
    // /// bound is inclusive; the upper bound is exclusive. Fails when
    // /// the upper bound is not greater than the lower bound.
    // fn gen_bigint_range(&mut self, lbound: &BigInt, ubound: &BigInt) -> BigInt;
}

impl<R: Rng> RandBigInt for R {
    fn gen_biguint(&mut self, bit_size: uint) -> BigUint {
        let mut state = RandState::new();
        let seed: c_ulong = self.gen();
        state.seed_ui(seed);
        let mpz = state.urandom_2exp(bit_size as c_ulong);
        BigUint{ data: mpz }
    }

    fn gen_biguint_below(&mut self, bound: &BigUint) -> BigUint {
        // FIXME Add assertion once is_zero is implemented
        // assert!(!bound.is_zero());

        let mut state = RandState::new();
        let seed: c_ulong = self.gen();
        state.seed_ui(seed);
        let mpz = state.urandom(&bound.data);
        BigUint{ data: mpz }
    }

    fn gen_biguint_range(&mut self, lbound: &BigUint, ubound: &BigUint) -> BigUint {
        // FIXME Add assertion once cmp::Ord is implemented
        // assert!(lbound < ubound);

        return *lbound + self.gen_biguint_below(&(*ubound - *lbound));
    }
}

#[cfg(test)]
mod test {
    use super::{BigUint, RandBigInt, ToBigUint};
    use std::{u32,u64};
    use std::rand::task_rng;

    #[test]
    fn test_to_and_fro() {
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();
        assert_eq!(two.to_str(), ~"2");
        let three = 3u.to_biguint().unwrap();
        assert_eq!(three.to_str(), ~"3");
    }

    #[test]
    fn test_printing() {
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();
        println!("{} == 2!", two); // Doesn't fail
    }

    #[test]
    fn test_add() {
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();
        let three: BigUint = FromPrimitive::from_uint(3).unwrap();

        assert_eq!(two.add(&three).to_str(), ~"5");
        assert_eq!((two + three).to_str(), ~"5");
    }

    #[test]
    fn test_simple_sub() {
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();
        let three: BigUint = FromPrimitive::from_uint(3).unwrap();

        assert_eq!(three.sub(&two).to_str(), ~"1");
        assert_eq!((three - two).to_str(), ~"1");
    }

    #[test]
    fn test_mul() {
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();
        let three: BigUint = FromPrimitive::from_uint(3).unwrap();

        assert_eq!(two.mul(&three).to_str(), ~"6");
        assert_eq!((two * three).to_str(), ~"6");
    }

    #[test]
    fn test_gen_biguint() {
        let mut rng = task_rng();
        let rand1 = rng.gen_biguint(128);
        let rand2 = rng.gen_biguint(128);
        assert!(rand1.to_str() != rand2.to_str());
    }

    #[test]
    fn test_gen_biguint_below() {
        let mut rng = task_rng();
        let max: BigUint = FromPrimitive::from_u64(u64::MAX).unwrap();
        let rand1 = rng.gen_biguint_below(&max);
        let rand2 = rng.gen_biguint_below(&max);
        assert!(rand1.to_str() != rand2.to_str());
    }

    #[test]
    fn test_gen_biguint_range() {
        let mut rng = task_rng();
        let min: BigUint = FromPrimitive::from_u32(u32::MAX).unwrap();
        let max: BigUint = FromPrimitive::from_u64(u64::MAX).unwrap();
        let rand1 = rng.gen_biguint_range(&min, &max);
        let rand2 = rng.gen_biguint_range(&min, &max);
        assert!(rand1.to_str() != rand2.to_str());
    }
}
<commit_msg>Implement more traits<commit_after>#[crate_id = "bignum#0.1.0-pre"];

#[comment = "Bignum library for Rust"];
#[crate_type = "rlib"];

#[feature(macro_rules)];

extern crate gmp;

use gmp::{Mpz, RandState};
use std::fmt;
use std::from_str::FromStr;
use std::num::{One, Zero, ToStrRadix};
use std::rand::Rng;
use std::libc::c_ulong;

#[deriving(Clone, Eq, Ord, TotalEq, TotalOrd, Zero)]
pub struct BigUint {
    data: Mpz
}

impl One for BigUint {
    fn one() -> BigUint {
        BigUint{ data: One::one() }
    }
}

impl FromPrimitive for BigUint {
    fn from_u64(other: u64) -> Option<BigUint> {
        let mpz: Option<Mpz> = FromPrimitive::from_u64(other);
        match mpz {
            Some(mpz) => Some(BigUint{ data: mpz }),
            None      => None
        }
    }

    fn from_i64(other: i64) -> Option<BigUint> {
        if other < 0 {
            None
        } else {
            let mpz: Option<Mpz> = FromPrimitive::from_i64(other);
            match mpz {
                Some(mpz) => Some(BigUint{ data: mpz }),
                None      => None
            }
        }
    }
}

pub trait ToBigUint {
    fn to_biguint(&self) -> Option<BigUint>;
}

macro_rules! impl_to_biguint(
    ($T:ty, $from_ty:path) => {
        impl ToBigUint for $T {
            fn to_biguint(&self) -> Option<BigUint> {
                $from_ty(*self)
            }
        }
    }
)

impl_to_biguint!(int,  FromPrimitive::from_int)
impl_to_biguint!(i8,   FromPrimitive::from_i8)
impl_to_biguint!(i16,  FromPrimitive::from_i16)
impl_to_biguint!(i32,  FromPrimitive::from_i32)
impl_to_biguint!(i64,  FromPrimitive::from_i64)
impl_to_biguint!(uint, FromPrimitive::from_uint)
impl_to_biguint!(u8,   FromPrimitive::from_u8)
impl_to_biguint!(u16,  FromPrimitive::from_u16)
impl_to_biguint!(u32,  FromPrimitive::from_u32)
impl_to_biguint!(u64,  FromPrimitive::from_u64)

impl FromStr for BigUint {
    fn from_str(s: &str) -> Option<BigUint> {
        let data: Option<Mpz> = FromStr::from_str(s);
        match data {
            Some(data) => Some(BigUint{ data: data }),
            None       => None
        }
    }
}

impl ToStrRadix for BigUint {
    fn to_str_radix(&self, radix: uint) -> ~str {
        self.data.to_str_radix(radix)
    }
}

impl fmt::Show for BigUint {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f.buf, "{}", self.to_str_radix(10))
    }
}

impl Add<BigUint, BigUint> for BigUint {
    fn add(&self, other: &BigUint) -> BigUint {
        BigUint{data: self.data.add(&other.data)}
    }
}

impl Sub<BigUint, BigUint> for BigUint {
    fn sub(&self, other: &BigUint) -> BigUint {
        BigUint{data: self.data.sub(&other.data)}
    }
}

impl Mul<BigUint, BigUint> for BigUint {
    fn mul(&self, other: &BigUint) -> BigUint {
        BigUint{data: self.data.mul(&other.data)}
    }
}

impl BigUint {
    pub fn is_odd(&self) -> bool {
        self.data.is_odd()
    }

    pub fn is_even(&self) -> bool {
        self.data.is_even()
    }
}

pub trait RandBigInt {
    /// Generate a random `BigUint` of the given bit size.
    fn gen_biguint(&mut self, bit_size: uint) -> BigUint;

    // /// Generate a random BigInt of the given bit size.
    // fn gen_bigint(&mut self, bit_size: uint) -> BigInt;

    /// Generate a random `BigUint` less than the given bound. Fails
    /// when the bound is zero.
    fn gen_biguint_below(&mut self, bound: &BigUint) -> BigUint;

    /// Generate a random `BigUint` within the given range. The lower
    /// bound is inclusive; the upper bound is exclusive. Fails when
    /// the upper bound is not greater than the lower bound.
    fn gen_biguint_range(&mut self, lbound: &BigUint, ubound: &BigUint) -> BigUint;

    // /// Generate a random `BigInt` within the given range. The lower
    // /// bound is inclusive; the upper bound is exclusive. Fails when
    // /// the upper bound is not greater than the lower bound.
    // fn gen_bigint_range(&mut self, lbound: &BigInt, ubound: &BigInt) -> BigInt;
}

impl<R: Rng> RandBigInt for R {
    fn gen_biguint(&mut self, bit_size: uint) -> BigUint {
        let mut state = RandState::new();
        let seed: c_ulong = self.gen();
        state.seed_ui(seed);
        let mpz = state.urandom_2exp(bit_size as c_ulong);
        BigUint{ data: mpz }
    }

    fn gen_biguint_below(&mut self, bound: &BigUint) -> BigUint {
        // FIXME Add assertion once is_zero is implemented
        // assert!(!bound.is_zero());

        let mut state = RandState::new();
        let seed: c_ulong = self.gen();
        state.seed_ui(seed);
        let mpz = state.urandom(&bound.data);
        BigUint{ data: mpz }
    }

    fn gen_biguint_range(&mut self, lbound: &BigUint, ubound: &BigUint) -> BigUint {
        // FIXME Add assertion once cmp::Ord is implemented
        // assert!(lbound < ubound);

        return *lbound + self.gen_biguint_below(&(*ubound - *lbound));
    }
}

#[cfg(test)]
mod test {
    use super::{BigUint, RandBigInt, ToBigUint};
    use std::{u32,u64};
    use std::from_str::FromStr;
    use std::num::{Zero, One};
    use std::rand::task_rng;

    #[test]
    fn test_clone() {
        let two = 2u.to_biguint().unwrap();
        let also_two = two.clone();
        assert_eq!(two.to_str(), also_two.to_str());
    }

    #[test]
    fn test_from_primitive() {
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();
        assert_eq!(two.to_str(), ~"2");
    }

    #[test]
    fn test_from_str() {
        let two: BigUint = FromStr::from_str("2").unwrap();
        assert_eq!(two.to_str(), ~"2");
    }

    #[test]
    fn test_to_biguint() {
        let three = 3u.to_biguint().unwrap();
        assert_eq!(three.to_str(), ~"3");
    }

    #[test]
    fn test_comparisons() {
        let two = 2u.to_biguint().unwrap();
        let also_two = 2u.to_biguint().unwrap();
        let three = 3u.to_biguint().unwrap();
        assert!(two == also_two);
        assert!(two >= also_two);
        assert!(two <= also_two);
        assert!(three > two);
        assert!(three >= two);
        assert!(two < three);
        assert!(two <= three);
    }

    #[test]
    fn test_zero_and_one() {
        let zero: BigUint = Zero::zero();
        let one: BigUint = One::one();
        assert_eq!(zero.to_str(), ~"0");
        assert_eq!(one.to_str(), ~"1");
    }

    #[test]
    fn test_printing() {
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();
        println!("{} == 2!", two); // Doesn't fail
    }

    #[test]
    fn test_add() {
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();
        let three: BigUint = FromPrimitive::from_uint(3).unwrap();

        assert_eq!(two.add(&three).to_str(), ~"5");
        assert_eq!((two + three).to_str(), ~"5");
    }

    #[test]
    fn test_simple_sub() {
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();
        let three: BigUint = FromPrimitive::from_uint(3).unwrap();

        assert_eq!(three.sub(&two).to_str(), ~"1");
        assert_eq!((three - two).to_str(), ~"1");
    }

    #[test]
    fn test_mul() {
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();
        let three: BigUint = FromPrimitive::from_uint(3).unwrap();

        assert_eq!(two.mul(&three).to_str(), ~"6");
        assert_eq!((two * three).to_str(), ~"6");
    }

    #[test]
    fn test_is_odd() {
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();
        let three: BigUint = FromPrimitive::from_uint(3).unwrap();

        assert!(!two.is_odd());
        assert!(three.is_odd());
    }

    #[test]
    fn test_is_even() {
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();
        let three: BigUint = FromPrimitive::from_uint(3).unwrap();

        assert!(two.is_even());
        assert!(!three.is_even());
    }

    #[test]
    fn test_rand_gen_biguint() {
        let mut rng = task_rng();
        let rand1 = rng.gen_biguint(128);
        let rand2 = rng.gen_biguint(128);
        assert!(rand1.to_str() != rand2.to_str());
    }

    #[test]
    fn test_rand_gen_biguint_below() {
        let mut rng = task_rng();
        let max: BigUint = FromPrimitive::from_u64(u64::MAX).unwrap();
        let rand1 = rng.gen_biguint_below(&max);
        let rand2 = rng.gen_biguint_below(&max);
        assert!(rand1.to_str() != rand2.to_str());
    }

    #[test]
    fn test_rand_gen_biguint_range() {
        let mut rng = task_rng();
        let min: BigUint = FromPrimitive::from_u32(u32::MAX).unwrap();
        let max: BigUint = FromPrimitive::from_u64(u64::MAX).unwrap();
        let rand1 = rng.gen_biguint_range(&min, &max);
        let rand2 = rng.gen_biguint_range(&min, &max);
        assert!(rand1.to_str() != rand2.to_str());
    }
}
<|endoftext|>


<commit_before>#![crate_type = "bin"]

extern crate cairo;
extern crate pango_sys;
extern crate gdk;
extern crate glib;
extern crate gtk;
extern crate gtk_sys;
extern crate libc;
extern crate sysinfo;

use sysinfo::*;

use gtk::prelude::*;
use gtk::{AboutDialog, Button, Dialog, Entry, MenuBar};

use std::cell::RefCell;
use std::collections::{HashMap, HashSet};
use std::ffi::CString;
use std::rc::Rc;
use std::time::Duration;
use std::thread::sleep;

use display_sysinfo::DisplaySysInfo;
use notebook::NoteBook;
use procs::{create_and_fill_model, Procs};

mod color;
mod display_sysinfo;
mod notebook;
mod process_dialog;
mod procs;
mod utils;

fn update_window(list: &gtk::ListStore, system: &Rc<RefCell<sysinfo::System>>,
                 info: &mut DisplaySysInfo) {
    let mut system = system.borrow_mut();
    system.refresh_all();
    info.update_ram_display(&system);
    info.update_process_display(&system);
    let entries: &HashMap<usize, Process> = system.get_process_list();
    let mut seen: HashSet<usize> = HashSet::new();

    if let Some(mut iter) = list.get_iter_first() {
        let mut valid = true;
        while valid {
            let pid = list.get_value(&iter, 0).get::<i64>().unwrap() as usize;
            if let Some(p) = entries.get(&(pid)) {
                list.set(&iter,
                         &[2, 3, 5],
                         &[&format!("{:.1}", p.cpu_usage), &p.memory, &p.cpu_usage]);
                valid = list.iter_next(&mut iter);
                seen.insert(pid);
            } else {
                valid = list.remove(&mut iter);
            }
        }
    }

    for (pid, pro) in entries.iter() {
        if !seen.contains(pid) {
            create_and_fill_model(list, pro.pid, &pro.cmd, &pro.name, pro.cpu_usage, pro.memory);
        }
    }
}

fn parse_quote(line: &str, quote: char) -> Vec<CString> {
    let args = line.split(quote).collect::<Vec<&str>>();
    let mut out_args = vec!();

    for (num, arg) in args.iter().enumerate() {
        if num != 1 {
            out_args.extend_from_slice(&parse_entry(arg));
        } else {
            out_args.push(CString::new(*arg).unwrap());
        }
    }
    out_args
}

// super simple parsing
fn parse_entry(line: &str) -> Vec<CString> {
    match (line.find('\''), line.find('"')) {
        (Some(x), Some(y)) => {
            if x < y {
                parse_quote(line, '\'')
            } else {
                parse_quote(line, '"')
            }
        }
        (Some(_), None) => parse_quote(line, '\''),
        (None, Some(_)) => parse_quote(line, '"'),
        (None, None) => line.split(' ').into_iter()
                                       .map(|s| CString::new(s).unwrap())
                                       .collect::<Vec<CString>>(),
    }
}

fn start_detached_process(line: &str) -> Option<i32> {
    let args = parse_entry(line);
    let command = args[0].clone();

    unsafe {
        let pid = libc::fork();

        if pid < 0 {
            libc::exit(2); // error
        } else if pid == 0 {
            libc::umask(0);
            if libc::setsid() < 0 {
                libc::exit(3);
            }
            libc::chdir("/".as_ptr() as *const i8);
            let stdin = libc::fdopen(0, "r".as_ptr() as *const i8);
            let stdout = libc::fdopen(1, "w".as_ptr() as *const i8);
            let stderr = libc::fdopen(2, "w".as_ptr() as *const i8);
            libc::freopen("/dev/null".as_ptr() as *const i8, "r".as_ptr() as *const i8, stdin);
            libc::freopen("/dev/null".as_ptr() as *const i8, "w".as_ptr() as *const i8, stdout);
            libc::freopen("/dev/null".as_ptr() as *const i8, "w".as_ptr() as *const i8, stderr);
            libc::execv(command.as_ptr(), args.into_iter()
                                              .map(|s| s.as_ptr())
                                              .collect::<Vec<*const libc::c_char>>().as_ptr());
            None
        } else {
            sleep(Duration::from_millis(100));
            if libc::kill(pid, 0) != 0 {
                let mut status = 0i32;
                // the process isn't running, getting error code
                libc::wait(&mut status as *mut i32);
                Some(status)
            } else {
                None
            }
        }
    }
}

fn main() {
    gtk::init().expect("GTK couldn't start normally");

    let window = gtk::Window::new(gtk::WindowType::Toplevel);
    let sys = sysinfo::System::new();
    let start_time = unsafe { sys.get_process(libc::getpid() as i64).unwrap().start_time };
    let running_since = Rc::new(RefCell::new(0));
    let sys = Rc::new(RefCell::new(sys));
    let mut note = NoteBook::new();
    let procs = Procs::new(sys.borrow().get_process_list(), &mut note);
    let current_pid = procs.current_pid.clone();
    let current_pid2 = procs.current_pid.clone();
    let sys1 = sys.clone();
    let sys2 = sys.clone();
    let sys3 = sys.clone();
    let info_button = procs.info_button.clone();

    window.set_title("Process viewer");
    window.set_position(gtk::WindowPosition::Center);

    window.connect_delete_event(|_, _| {
        gtk::main_quit();
        Inhibit(false)
    });

    sys.borrow_mut().refresh_all();
    procs.kill_button.connect_clicked(move |_| {
        let sys = sys1.borrow();
        if let Some(process) = current_pid.get().and_then(|pid| sys.get_process(pid)) {
            process.kill(Signal::Kill);
        }
    });

    let mut display_tab = DisplaySysInfo::new(sys.clone(), &mut note, &window);

    let v_box = gtk::Box::new(gtk::Orientation::Vertical, 0);
    let menu_bar = MenuBar::new();
    let menu = gtk::Menu::new();
    let more_menu = gtk::Menu::new();
    let file = gtk::MenuItem::new_with_label("File");
    let new_task = gtk::MenuItem::new_with_label("Launch new executable");
    let quit = gtk::MenuItem::new_with_label("Quit");
    let more = gtk::MenuItem::new_with_label("?");
    let about = gtk::MenuItem::new_with_label("About");

    menu.append(&new_task);
    menu.append(&quit);
    file.set_submenu(Some(&menu));
    menu_bar.append(&file);
    more_menu.append(&about);
    more.set_submenu(Some(&more_menu));
    menu_bar.append(&more);

    v_box.pack_start(&menu_bar, false, false, 0);
    v_box.pack_start(&note.notebook, true, true, 0);

    window.add(&v_box);
    window.show_all();
    let window1 = window.clone();
    let window2 = window.clone();

    let list_store = procs.list_store.clone();
    let run1 = running_since.clone();
    let run2 = running_since.clone();
    gtk::timeout_add(1000, move || {
        *run1.borrow_mut() += 1;
        // first part, deactivate sorting
        let sorted = list_store.get_sort_column_id();
        list_store.set_unsorted();

        // we update the tree view
        update_window(&list_store, &sys, &mut display_tab);

        // we re-enable the sorting
        if let Some((col, order)) = sorted {
            list_store.set_sort_column_id(col, order);
        }
        window1.queue_draw();
        glib::Continue(true)
    });
    info_button.connect_clicked(move |_| {
        let sys = sys2.borrow();
        if let Some(process) = current_pid2.get().and_then(|pid| sys.get_process(pid)) {
            process_dialog::create_process_dialog(&process, &window2, start_time,
                                                  run2.borrow().clone());
        }
    });
    let window2 = window.clone();
    procs.left_tree.connect_row_activated(move |tree_view, path, _| {
        let model = tree_view.get_model().unwrap();
        let iter = model.get_iter(path).unwrap();
        let sys = sys3.borrow();
        if let Some(process) = sys.get_process(model.get_value(&iter, 0).get().unwrap()) {
            process_dialog::create_process_dialog(&process, &window2, start_time,
                                                  running_since.borrow().clone());
        }
    });

    quit.connect_activate(|_| {
        gtk::main_quit();
    });
    let window3 = window.clone();
    about.connect_activate(move |_| {
        let p = AboutDialog::new();
        p.set_authors(&["Guillaume Gomez"]);
        p.set_website_label(Some("my website"));
        p.set_website(Some("http://guillaume-gomez.fr/"));
        p.set_comments(Some("A process viewer GUI wrote with gtk-rs"));
        p.set_copyright(Some("This is under MIT license"));
        p.set_transient_for(Some(&window3));
        p.show_all();
    });
    new_task.connect_activate(move |_| {
        let d = Dialog::new();
        d.set_title("Launch new executable");
        let content_area = d.get_content_area();
        let input = Entry::new();
        let run = Button::new_with_label("Run");
        let cancel = Button::new_with_label("Cancel");
        let v_box = gtk::Box::new(gtk::Orientation::Vertical, 0);
        let h_box = gtk::Box::new(gtk::Orientation::Horizontal, 0);
        h_box.pack_start(&run, true, true, 0);
        h_box.pack_start(&cancel, true, true, 0);
        v_box.pack_start(&input, true, true, 0);
        v_box.pack_start(&h_box, true, true, 0);
        content_area.add(&v_box);
        d.set_transient_for(Some(&window));
        d.set_size_request(400, 70);
        d.show_all();

        run.set_sensitive(false);
        let run2 = run.clone();
        let input2 = input.clone();
        input.connect_changed(move |_| {
            match input2.get_text() {
                Some(ref x) if x.len() > 0 => run2.set_sensitive(true),
                _ => run2.set_sensitive(false),
            }
        });
        cancel.connect_clicked(move |_| {
            d.destroy();
        });
        run.connect_clicked(move |_| {
            if let Some(text) = input.get_text() {
                start_detached_process(&text);
            }
        });
    });

    gtk::main();
}
<commit_msg>Add activate signal on run command input<commit_after>#![crate_type = "bin"]

extern crate cairo;
extern crate pango_sys;
extern crate gdk;
extern crate glib;
extern crate gtk;
extern crate gtk_sys;
extern crate libc;
extern crate sysinfo;

use sysinfo::*;

use gtk::prelude::*;
use gtk::{AboutDialog, Button, Dialog, Entry, MenuBar, MessageDialog};

use std::cell::RefCell;
use std::collections::{HashMap, HashSet};
use std::ffi::CString;
use std::rc::Rc;
use std::time::Duration;
use std::thread::sleep;

use display_sysinfo::DisplaySysInfo;
use notebook::NoteBook;
use procs::{create_and_fill_model, Procs};

mod color;
mod display_sysinfo;
mod notebook;
mod process_dialog;
mod procs;
mod utils;

fn update_window(list: &gtk::ListStore, system: &Rc<RefCell<sysinfo::System>>,
                 info: &mut DisplaySysInfo) {
    let mut system = system.borrow_mut();
    system.refresh_all();
    info.update_ram_display(&system);
    info.update_process_display(&system);
    let entries: &HashMap<usize, Process> = system.get_process_list();
    let mut seen: HashSet<usize> = HashSet::new();

    if let Some(mut iter) = list.get_iter_first() {
        let mut valid = true;
        while valid {
            let pid = list.get_value(&iter, 0).get::<i64>().unwrap() as usize;
            if let Some(p) = entries.get(&(pid)) {
                list.set(&iter,
                         &[2, 3, 5],
                         &[&format!("{:.1}", p.cpu_usage), &p.memory, &p.cpu_usage]);
                valid = list.iter_next(&mut iter);
                seen.insert(pid);
            } else {
                valid = list.remove(&mut iter);
            }
        }
    }

    for (pid, pro) in entries.iter() {
        if !seen.contains(pid) {
            create_and_fill_model(list, pro.pid, &pro.cmd, &pro.name, pro.cpu_usage, pro.memory);
        }
    }
}

fn parse_quote(line: &str, quote: char) -> Vec<CString> {
    let args = line.split(quote).collect::<Vec<&str>>();
    let mut out_args = vec!();

    for (num, arg) in args.iter().enumerate() {
        if num != 1 {
            out_args.extend_from_slice(&parse_entry(arg));
        } else {
            out_args.push(CString::new(*arg).unwrap());
        }
    }
    out_args
}

// super simple parsing
fn parse_entry(line: &str) -> Vec<CString> {
    match (line.find('\''), line.find('"')) {
        (Some(x), Some(y)) => {
            if x < y {
                parse_quote(line, '\'')
            } else {
                parse_quote(line, '"')
            }
        }
        (Some(_), None) => parse_quote(line, '\''),
        (None, Some(_)) => parse_quote(line, '"'),
        (None, None) => line.split(' ').into_iter()
                                       .map(|s| CString::new(s).unwrap())
                                       .collect::<Vec<CString>>(),
    }
}

fn start_detached_process(line: &str) -> Option<i32> {
    let args = parse_entry(line);
    let command = args[0].clone();
    let mut args = args.into_iter()
                       .map(|s| s.as_ptr())
                       .collect::<Vec<*const libc::c_char>>();
    args.push(std::ptr::null());

    unsafe {
        let pid = libc::fork();

        if pid < 0 {
            libc::exit(2); // error
        } else if pid == 0 {
            libc::umask(0);
            if libc::setsid() < 0 {
                libc::exit(3);
            }
            libc::chdir("/".as_ptr() as *const i8);
            /*let stdin = libc::fdopen(0, "r".as_ptr() as *const i8);
            let stdout = libc::fdopen(1, "w".as_ptr() as *const i8);
            let stderr = libc::fdopen(2, "w".as_ptr() as *const i8);
            libc::freopen("/dev/null".as_ptr() as *const i8, "r".as_ptr() as *const i8, stdin);
            libc::freopen("/dev/null".as_ptr() as *const i8, "w".as_ptr() as *const i8, stdout);
            libc::freopen("/dev/null".as_ptr() as *const i8, "w".as_ptr() as *const i8, stderr);*/
            libc::execv(command.as_ptr(), args.as_ptr());
            None
        } else {
            sleep(Duration::from_millis(100));
            if libc::kill(pid, 0) != 0 {
                let mut status = 0i32;
                // the process isn't running, getting error code
                libc::wait(&mut status as *mut i32);
                Some(status)
            } else {
                None
            }
        }
    }
}

fn run_command(input: &Entry, window: &gtk::Window, d: &Dialog) {
    if let Some(text) = input.get_text() {
        let x = if let Some(x) = start_detached_process(&text) {
            x
        } else {
            0
        };
        d.destroy();
        let m = MessageDialog::new(Some(window),
                                   gtk::DIALOG_DESTROY_WITH_PARENT,
                                   gtk::MessageType::Info,
                                   gtk::ButtonsType::Ok,
                                   &format!("The command exited with {} code", x));
        m.show_all();
    }
}

fn main() {
    gtk::init().expect("GTK couldn't start normally");

    let window = gtk::Window::new(gtk::WindowType::Toplevel);
    let sys = sysinfo::System::new();
    let start_time = unsafe { sys.get_process(libc::getpid() as i64).unwrap().start_time };
    let running_since = Rc::new(RefCell::new(0));
    let sys = Rc::new(RefCell::new(sys));
    let mut note = NoteBook::new();
    let procs = Procs::new(sys.borrow().get_process_list(), &mut note);
    let current_pid = procs.current_pid.clone();
    let current_pid2 = procs.current_pid.clone();
    let sys1 = sys.clone();
    let sys2 = sys.clone();
    let sys3 = sys.clone();
    let info_button = procs.info_button.clone();

    window.set_title("Process viewer");
    window.set_position(gtk::WindowPosition::Center);

    window.connect_delete_event(|_, _| {
        gtk::main_quit();
        Inhibit(false)
    });

    sys.borrow_mut().refresh_all();
    procs.kill_button.connect_clicked(move |_| {
        let sys = sys1.borrow();
        if let Some(process) = current_pid.get().and_then(|pid| sys.get_process(pid)) {
            process.kill(Signal::Kill);
        }
    });

    let mut display_tab = DisplaySysInfo::new(sys.clone(), &mut note, &window);

    let v_box = gtk::Box::new(gtk::Orientation::Vertical, 0);
    let menu_bar = MenuBar::new();
    let menu = gtk::Menu::new();
    let more_menu = gtk::Menu::new();
    let file = gtk::MenuItem::new_with_label("File");
    let new_task = gtk::MenuItem::new_with_label("Launch new executable");
    let quit = gtk::MenuItem::new_with_label("Quit");
    let more = gtk::MenuItem::new_with_label("?");
    let about = gtk::MenuItem::new_with_label("About");

    menu.append(&new_task);
    menu.append(&quit);
    file.set_submenu(Some(&menu));
    menu_bar.append(&file);
    more_menu.append(&about);
    more.set_submenu(Some(&more_menu));
    menu_bar.append(&more);

    v_box.pack_start(&menu_bar, false, false, 0);
    v_box.pack_start(&note.notebook, true, true, 0);

    window.add(&v_box);
    window.show_all();
    let window1 = window.clone();
    let window2 = window.clone();

    let list_store = procs.list_store.clone();
    let run1 = running_since.clone();
    let run2 = running_since.clone();
    gtk::timeout_add(1000, move || {
        *run1.borrow_mut() += 1;
        // first part, deactivate sorting
        let sorted = list_store.get_sort_column_id();
        list_store.set_unsorted();

        // we update the tree view
        update_window(&list_store, &sys, &mut display_tab);

        // we re-enable the sorting
        if let Some((col, order)) = sorted {
            list_store.set_sort_column_id(col, order);
        }
        window1.queue_draw();
        glib::Continue(true)
    });
    info_button.connect_clicked(move |_| {
        let sys = sys2.borrow();
        if let Some(process) = current_pid2.get().and_then(|pid| sys.get_process(pid)) {
            process_dialog::create_process_dialog(&process, &window2, start_time,
                                                  run2.borrow().clone());
        }
    });
    let window2 = window.clone();
    procs.left_tree.connect_row_activated(move |tree_view, path, _| {
        let model = tree_view.get_model().unwrap();
        let iter = model.get_iter(path).unwrap();
        let sys = sys3.borrow();
        if let Some(process) = sys.get_process(model.get_value(&iter, 0).get().unwrap()) {
            process_dialog::create_process_dialog(&process, &window2, start_time,
                                                  running_since.borrow().clone());
        }
    });

    quit.connect_activate(|_| {
        gtk::main_quit();
    });
    let window3 = window.clone();
    about.connect_activate(move |_| {
        let p = AboutDialog::new();
        p.set_authors(&["Guillaume Gomez"]);
        p.set_website_label(Some("my website"));
        p.set_website(Some("http://guillaume-gomez.fr/"));
        p.set_comments(Some("A process viewer GUI wrote with gtk-rs"));
        p.set_copyright(Some("This is under MIT license"));
        p.set_transient_for(Some(&window3));
        p.show_all();
    });
    new_task.connect_activate(move |_| {
        let d = Dialog::new();
        d.set_title("Launch new executable");
        let content_area = d.get_content_area();
        let input = Entry::new();
        let run = Button::new_with_label("Run");
        let cancel = Button::new_with_label("Cancel");
        let v_box = gtk::Box::new(gtk::Orientation::Vertical, 0);
        let h_box = gtk::Box::new(gtk::Orientation::Horizontal, 0);
        h_box.pack_start(&run, true, true, 0);
        h_box.pack_start(&cancel, true, true, 0);
        v_box.pack_start(&input, true, true, 0);
        v_box.pack_start(&h_box, true, true, 0);
        content_area.add(&v_box);
        let window2 = window.clone();
        d.set_transient_for(Some(&window2));
        d.set_size_request(400, 70);
        d.show_all();

        run.set_sensitive(false);
        let run2 = run.clone();
        let input2 = input.clone();
        input.connect_changed(move |_| {
            match input2.get_text() {
                Some(ref x) if x.len() > 0 => run2.set_sensitive(true),
                _ => run2.set_sensitive(false),
            }
        });
        let d2 = d.clone();
        cancel.connect_clicked(move |_| {
            d2.destroy();
        });
        let window3 = window.clone();
        let input3 = input.clone();
        let d3 = d.clone();
        input.connect_activate(move |_| {
            run_command(&input3, &window3, &d3);
        });
        let window4 = window.clone();
        run.connect_clicked(move |_| {
            run_command(&input, &window4, &d);
        });
    });

    gtk::main();
}
<|endoftext|>


<commit_before>use yaml_rust::Yaml;
use command::{Command, Session, SendKeys, Split, Layout, Window, Window2, Attach, Panes};

#[cfg(test)] use yaml_rust::{YamlLoader};

pub fn main(yaml_string: &Vec<Yaml>, project_name: String) -> Vec<Command> {
    let mut commands: Vec<Command> = vec!();

    for doc in yaml_string {
        let root = match doc["root"].as_str() {
            Some(x) => Some(x.to_string()),
            None    => None
        };

        let (first_window, windows) = doc["windows"].as_vec().expect("No Windows have been defined.").split_at(1);

        match &first_window[0] {
             &Yaml::Hash(ref h)  => {
                 for (k, v) in h {
                     if v.as_hash().is_some() {
                         commands.push(Command::Session(Session{name: project_name.clone(), window_name: k.as_str().unwrap().to_string(), root: root.clone()}));
                         commands.append(&mut pane_matcher(project_name.clone(), v, root.clone(), k.as_str().unwrap().to_string()));
                     } else {
                         commands.push(Command::Session(Session{name: project_name.clone(), window_name: k.as_str().unwrap().to_string(), root: root.clone()}));
                     }
                 }
             },
             &Yaml::String(ref s) => {
                 commands.push(Command::Session(Session{name: project_name.clone(), window_name: s.clone(), root: root.clone()}))
             },
             &Yaml::Integer(ref s) => {
                 commands.push(Command::Session(Session{name: project_name.clone(), window_name: s.to_string(), root: root.clone()}))
             },
             _ => panic!("Muxed config file formatting isn't recognized.")
        };

        for (i,window) in windows.iter().enumerate() {
            match window {
                &Yaml::Hash(ref h)  => {
                    for (k, v) in h {
                        if v.as_hash().is_some() {
                            commands.push(Command::Window(Window{value: k.as_str().unwrap().to_string(), root: root.clone(), exec: "".to_string()}));
                            commands.append(&mut pane_matcher(project_name.clone(), v, root.clone(), k.as_str().unwrap().to_string()));
                        } else {
                            commands.push(Command::Window2(Window2{session_name: project_name.clone(), name: k.as_str().unwrap().to_string(), root: root.clone()}));
                            commands.push(Command::SendKeys(SendKeys{target: format!("{}:{}", project_name, k.as_str().unwrap().to_string()).to_string(), exec: v.as_str().expect("Bad exec command").to_string()}));
                        }
                    }
                },
                &Yaml::String(ref s) => {
                    commands.push(Command::Window2(Window2{session_name: project_name.clone(), name: s.clone(), root: root.clone()}))
                },
                &Yaml::Integer(ref s) => {
                    commands.push(Command::Window2(Window2{session_name: project_name.clone(), name: s.to_string(), root: root.clone()}))
                },
                _ => panic!("Muxed config file formatting isn't recognized.")
            };
        };
    };

    commands
}

fn pane_matcher(session: String, panes: &Yaml, root: Option<String>, window: String) -> Vec<Command> {
    let mut commands = vec!();
    let panes2 = panes["panes"].as_vec().expect("Something is wrong with panes.");

    if panes["layout"].as_str().is_some() {
        let layout = panes["layout"].as_str().unwrap().to_string();
        commands.push(Command::Layout(Layout{panes: panes2.len(), target: format!("{}:{}", session, window).to_string(), layout: panes["layout"].as_str().expect("Bad layout").to_string()}));
    };

    for (i, pane) in panes2.iter().enumerate() {
        commands.push(Command::SendKeys(SendKeys{target: format!("{}:{}.{}", session, window, i).to_string(), exec: pane.as_str().expect("Bad exec command").to_string()}));
    };

    commands
}

#[test]
pub fn windows_as_array() {
    let s = "---
windows: ['cargo', 'vim', 'git']
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let commands = main(&yaml, "muxed".to_string());
    assert_eq!(commands.len(), 3)
}

#[test]
pub fn windows_with_integer_names() {
    let s = "---
windows: [1, 'vim', 3]
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let commands = main(&yaml, "muxed".to_string());
    assert_eq!(commands.len(), 3)
}

#[test]
pub fn windows_as_list() {
    let s = "---
windows:
  - cargo: ''
  - vim: ''
  - git: ''
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let commands = main(&yaml, "muxed".to_string());
    assert_eq!(commands.len(), 3)
}

#[test]
pub fn root_command() {
    let s = "---
root: '~/.muxed'
windows:
  - cargo: ''
  - vim: ''
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let commands = main(&yaml, "muxed".to_string());

    let first_window: Option<Window> = match commands[0].clone() {
        Command::Window(w) => Some(w),
        _                  => None
    };

    assert_eq!(first_window.unwrap().root.unwrap(), "~/.muxed".to_string())
}

#[test]
pub fn panes_array() {
    let s = "---
windows:
  - editor:
      layout: 'main-vertical'
      panes: ['vim', 'guard']
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let commands = main(&yaml, "muxed".to_string());
    assert_eq!(commands.len(), 2)
}

#[test]
pub fn panes_command_exists() {
    let s = "---
windows:
  - editor:
      layout: 'main-vertical'
      panes: ['vim', 'guard']
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let commands = main(&yaml, "muxed".to_string());

    let pane_command: Option<Panes> = match commands[1].clone() {
        Command::Panes(w) => Some(w),
        _                 => None
    };

    assert!(pane_command.is_some())
}

#[test]
pub fn panes_command_execs_array() {
    let s = "---
windows:
  - editor:
      layout: 'main-vertical'
      panes: ['vim', 'guard']
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let commands = main(&yaml, "muxed".to_string());

    let pane_command: Option<Panes> = match commands[1].clone() {
        Command::Panes(w) => Some(w),
        _                 => None
    };

    assert_eq!(pane_command.clone().unwrap().exec[0], "vim");
    assert_eq!(pane_command.clone().unwrap().exec[1], "guard")
}
<commit_msg>Finish it off by attaching.<commit_after>use yaml_rust::Yaml;
use command::{Command, Session, SendKeys, Split, Layout, Window, Window2, Attach, Panes};

#[cfg(test)] use yaml_rust::{YamlLoader};

pub fn main(yaml_string: &Vec<Yaml>, project_name: String) -> Vec<Command> {
    let mut commands: Vec<Command> = vec!();

    for doc in yaml_string {
        let root = match doc["root"].as_str() {
            Some(x) => Some(x.to_string()),
            None    => None
        };

        let (first_window, windows) = doc["windows"].as_vec().expect("No Windows have been defined.").split_at(1);

        match &first_window[0] {
             &Yaml::Hash(ref h)  => {
                 for (k, v) in h {
                     if v.as_hash().is_some() {
                         commands.push(Command::Session(Session{name: project_name.clone(), window_name: k.as_str().unwrap().to_string(), root: root.clone()}));
                         commands.append(&mut pane_matcher(project_name.clone(), v, root.clone(), k.as_str().unwrap().to_string()));
                     } else {
                         commands.push(Command::Session(Session{name: project_name.clone(), window_name: k.as_str().unwrap().to_string(), root: root.clone()}));
                     }
                 }
             },
             &Yaml::String(ref s) => {
                 commands.push(Command::Session(Session{name: project_name.clone(), window_name: s.clone(), root: root.clone()}))
             },
             &Yaml::Integer(ref s) => {
                 commands.push(Command::Session(Session{name: project_name.clone(), window_name: s.to_string(), root: root.clone()}))
             },
             _ => panic!("Muxed config file formatting isn't recognized.")
        };

        for (i,window) in windows.iter().enumerate() {
            match window {
                &Yaml::Hash(ref h)  => {
                    for (k, v) in h {
                        if v.as_hash().is_some() {
                            commands.push(Command::Window(Window{value: k.as_str().unwrap().to_string(), root: root.clone(), exec: "".to_string()}));
                            commands.append(&mut pane_matcher(project_name.clone(), v, root.clone(), k.as_str().unwrap().to_string()));
                        } else {
                            commands.push(Command::Window2(Window2{session_name: project_name.clone(), name: k.as_str().unwrap().to_string(), root: root.clone()}));
                            commands.push(Command::SendKeys(SendKeys{target: format!("{}:{}", project_name, k.as_str().unwrap().to_string()).to_string(), exec: v.as_str().expect("Bad exec command").to_string()}));
                        }
                    }
                },
                &Yaml::String(ref s) => {
                    commands.push(Command::Window2(Window2{session_name: project_name.clone(), name: s.clone(), root: root.clone()}))
                },
                &Yaml::Integer(ref s) => {
                    commands.push(Command::Window2(Window2{session_name: project_name.clone(), name: s.to_string(), root: root.clone()}))
                },
                _ => panic!("Muxed config file formatting isn't recognized.")
            };
        };
    };
    commands.push(Command::Attach(Attach{name: project_name.clone()}));
    commands
}

fn pane_matcher(session: String, panes: &Yaml, root: Option<String>, window: String) -> Vec<Command> {
    let mut commands = vec!();
    let panes2 = panes["panes"].as_vec().expect("Something is wrong with panes.");

    if panes["layout"].as_str().is_some() {
        let layout = panes["layout"].as_str().unwrap().to_string();
        commands.push(Command::Layout(Layout{panes: panes2.len(), target: format!("{}:{}", session, window).to_string(), layout: panes["layout"].as_str().expect("Bad layout").to_string()}));
    };

    for (i, pane) in panes2.iter().enumerate() {
        commands.push(Command::SendKeys(SendKeys{target: format!("{}:{}.{}", session, window, i).to_string(), exec: pane.as_str().expect("Bad exec command").to_string()}));
    };

    commands
}

#[test]
pub fn windows_as_array() {
    let s = "---
windows: ['cargo', 'vim', 'git']
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let commands = main(&yaml, "muxed".to_string());
    assert_eq!(commands.len(), 3)
}

#[test]
pub fn windows_with_integer_names() {
    let s = "---
windows: [1, 'vim', 3]
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let commands = main(&yaml, "muxed".to_string());
    assert_eq!(commands.len(), 3)
}

#[test]
pub fn windows_as_list() {
    let s = "---
windows:
  - cargo: ''
  - vim: ''
  - git: ''
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let commands = main(&yaml, "muxed".to_string());
    assert_eq!(commands.len(), 3)
}

#[test]
pub fn root_command() {
    let s = "---
root: '~/.muxed'
windows:
  - cargo: ''
  - vim: ''
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let commands = main(&yaml, "muxed".to_string());

    let first_window: Option<Window> = match commands[0].clone() {
        Command::Window(w) => Some(w),
        _                  => None
    };

    assert_eq!(first_window.unwrap().root.unwrap(), "~/.muxed".to_string())
}

#[test]
pub fn panes_array() {
    let s = "---
windows:
  - editor:
      layout: 'main-vertical'
      panes: ['vim', 'guard']
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let commands = main(&yaml, "muxed".to_string());
    assert_eq!(commands.len(), 2)
}

#[test]
pub fn panes_command_exists() {
    let s = "---
windows:
  - editor:
      layout: 'main-vertical'
      panes: ['vim', 'guard']
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let commands = main(&yaml, "muxed".to_string());

    let pane_command: Option<Panes> = match commands[1].clone() {
        Command::Panes(w) => Some(w),
        _                 => None
    };

    assert!(pane_command.is_some())
}

#[test]
pub fn panes_command_execs_array() {
    let s = "---
windows:
  - editor:
      layout: 'main-vertical'
      panes: ['vim', 'guard']
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let commands = main(&yaml, "muxed".to_string());

    let pane_command: Option<Panes> = match commands[1].clone() {
        Command::Panes(w) => Some(w),
        _                 => None
    };

    assert_eq!(pane_command.clone().unwrap().exec[0], "vim");
    assert_eq!(pane_command.clone().unwrap().exec[1], "guard")
}
<|endoftext|>


<commit_before>//! The YAML parser. Here is where we convert the yaml in to commands to be
/// processed later.

use yaml_rust::Yaml;
use command::{Command, Session, SendKeys, Split, Layout, Window, Attach, KillWindow};
use rand::random;

#[cfg(test)] use yaml_rust::{YamlLoader};

/// Here was pass in the parsed yaml and project name. The purpose of this main
/// loop is to build the stack of commands that are run to setup a users tmux
/// session.
///
/// yaml_string: The parsed yaml from the config file.
/// project_name: The name of the project.
pub fn main(yaml_string: &Vec<Yaml>, project_name: &String, daemonize: bool) -> Result<Vec<Command>, String> {
    let mut commands: Vec<Command> = vec!();
    let tmp_window_name = format!("muxed_first_window_{}", random::<u16>());

    // The initial session command. Contains the tmp_window_name to be closed
    // before attaching.
    commands.push(Command::Session(Session{name: project_name.clone(), tmp_window_name: tmp_window_name.clone()}));

    // There should only be one doc but it's a vec so loop it.
    for doc in yaml_string {
        let root = match doc["root"].as_str() {
            Some(x) => Some(x.to_string()),
            None    => None
        };

        let windows = doc["windows"].as_vec().expect("No Windows have been defined.");
        for (i, window) in windows.iter().enumerate() {
            match window {
                &Yaml::Hash(ref h)  => {
                    for (k, v) in h {
                        if v.as_hash().is_some() {
                            commands.push(Command::Window(Window{session_name: format!("{}:{}", project_name, i+1), name: k.as_str().unwrap().to_string(), root: root.clone()}));
                            commands.append(&mut try!(pane_matcher(&project_name, v, &root, k.as_str().unwrap().to_string())));
                        } else {
                            commands.push(Command::Window(Window{session_name: format!("{}:{}", project_name, i+1), name: k.as_str().unwrap().to_string(), root: root.clone()}));
                            commands.push(Command::SendKeys(SendKeys{target: format!("{}:{}", project_name, k.as_str().unwrap().to_string()).to_string(), exec: v.as_str().expect("Bad exec command").to_string()}));
                        }
                    }
                },
                &Yaml::String(ref s) => {
                    commands.push(Command::Window(Window{session_name: format!("{}:{}", project_name, i+1), name: s.clone(), root: root.clone()}))
                },
                &Yaml::Integer(ref s) => {
                    commands.push(Command::Window(Window{session_name: format!("{}:{}", project_name, i+1), name: s.to_string(), root: root.clone()}))
                },
                _ => panic!("Muxed config file formatting isn't recognized.")
            };
        };
    };

    commands.push(Command::KillWindow(KillWindow{name: format!("{}:{}", project_name, tmp_window_name)}));

    if !daemonize { commands.push(Command::Attach(Attach{name: project_name.clone()})) };
    Ok(commands)
}

/// Pane matcher is for breaking apart the panes. Splitting windows when needed
/// and executing commands as needed.
fn pane_matcher(session: &String, window: &Yaml, root: &Option<String>, window_name: String) -> Result<Vec<Command>, String> {
    let mut commands = vec!();
    let panes = window["panes"].as_vec().expect("Something is wrong with panes.");

    for (i, pane) in panes.iter().enumerate() {
        // For every pane, we need one less split.
        // ex. An existing window to become 2 panes, needs 1 split.
        if i < (panes.len()-1) {
            commands.push(Command::Split(Split{target: format!("{}:{}.{}", session, window_name, i).to_string(), root: root.clone()}));
        };
        // Execute given commands in each new pane after all splits are
        // complete.
        commands.push(Command::SendKeys(SendKeys{target: format!("{}:{}.{}", session, window_name, i).to_string(), exec: pane.as_str().expect("Bad exec command").to_string()}));
    };

    // After all panes are split select the layout for the window
    if window["layout"].as_str().is_some() {
        let layout = window["layout"].as_str().expect("Bad layout").to_string();
        commands.push(Command::Layout(Layout{target: format!("{}:{}", session, window_name).to_string(), layout: layout}));
    };

    Ok(commands)
}

#[test]
pub fn windows_defined_as_array_has_6_commands() {
    let s = "---
windows: ['cargo', 'vim', 'git']
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    assert_eq!(main(&yaml, &"muxed".to_string()).unwrap().len(), 6)
}

#[test]
pub fn windows_defined_as_array_has_1_session() {
    let s = "---
windows: ['cargo', 'vim', 'git']
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string()).unwrap().into_iter().filter(|x| match x {
        &Command::Session(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 1)
}

#[test]
pub fn windows_defined_as_array_has_3_windows() {
    let s = "---
windows: ['cargo', 'vim', 'git']
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string()).unwrap().into_iter().filter(|x| match x {
        &Command::Window(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 3)
}

#[test]
pub fn windows_defined_as_array_has_1_attach() {
    let s = "---
windows: ['cargo', 'vim', 'git']
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string()).unwrap().into_iter().filter(|x| match x {
        &Command::Attach(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 1)
}

#[test]
pub fn windows_with_integer_names() {
    let s = "---
windows: [1, 'vim', 3]
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    assert_eq!(main(&yaml, &"muxed".to_string()).unwrap().len(), 6)
}

#[test]
pub fn windows_as_list() {
    let s = "---
windows:
  - cargo: ''
  - vim: ''
  - git: ''
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let commands = main(&yaml, &"muxed".to_string()).unwrap();
    assert_eq!(commands.len(), 9)
}

#[test]
pub fn panes_array_has_8_commands() {
    let s = "---
windows:
  - editor:
      layout: 'main-vertical'
      panes: ['vim', 'guard']
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let commands = main(&yaml, &"muxed".to_string()).unwrap();
    assert_eq!(commands.len(), 8)
}

#[test]
pub fn panes_array_has_1_split() {
    let s = "---
windows:
  - editor:
      layout: 'main-vertical'
      panes: ['vim', 'guard']
";

    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string()).unwrap().into_iter().filter(|x| match x {
        &Command::Split(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 1)
}

#[test]
pub fn panes_array_has_1_layout() {
    let s = "---
windows:
  - editor:
      layout: 'main-vertical'
      panes: ['vim', 'guard']
";

    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string()).unwrap().into_iter().filter(|x| match x {
        &Command::Layout(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 1)
}

#[test]
pub fn panes_array_has_1_window() {
    let s = "---
windows:
  - editor:
      layout: 'main-vertical'
      panes: ['vim', 'guard']
";

    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string()).unwrap().into_iter().filter(|x| match x {
        &Command::Window(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 1)
}

#[test]
pub fn panes_array_has_1_session() {
    let s = "---
windows:
  - editor:
      layout: 'main-vertical'
      panes: ['vim', 'guard']
";

    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string()).unwrap().into_iter().filter(|x| match x {
        &Command::Session(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 1)
}
<commit_msg>Fix broken tests<commit_after>//! The YAML parser. Here is where we convert the yaml in to commands to be
/// processed later.

use yaml_rust::Yaml;
use command::{Command, Session, SendKeys, Split, Layout, Window, Attach, KillWindow};
use rand::random;

#[cfg(test)] use yaml_rust::{YamlLoader};

/// Here was pass in the parsed yaml and project name. The purpose of this main
/// loop is to build the stack of commands that are run to setup a users tmux
/// session.
///
/// yaml_string: The parsed yaml from the config file.
/// project_name: The name of the project.
pub fn main(yaml_string: &Vec<Yaml>, project_name: &String, daemonize: bool) -> Result<Vec<Command>, String> {
    let mut commands: Vec<Command> = vec!();
    let tmp_window_name = format!("muxed_first_window_{}", random::<u16>());

    // The initial session command. Contains the tmp_window_name to be closed
    // before attaching.
    commands.push(Command::Session(Session{name: project_name.clone(), tmp_window_name: tmp_window_name.clone()}));

    // There should only be one doc but it's a vec so loop it.
    for doc in yaml_string {
        let root = match doc["root"].as_str() {
            Some(x) => Some(x.to_string()),
            None    => None
        };

        let windows = doc["windows"].as_vec().expect("No Windows have been defined.");
        for (i, window) in windows.iter().enumerate() {
            match window {
                &Yaml::Hash(ref h)  => {
                    for (k, v) in h {
                        if v.as_hash().is_some() {
                            commands.push(Command::Window(Window{session_name: format!("{}:{}", project_name, i+1), name: k.as_str().unwrap().to_string(), root: root.clone()}));
                            commands.append(&mut try!(pane_matcher(&project_name, v, &root, k.as_str().unwrap().to_string())));
                        } else {
                            commands.push(Command::Window(Window{session_name: format!("{}:{}", project_name, i+1), name: k.as_str().unwrap().to_string(), root: root.clone()}));
                            commands.push(Command::SendKeys(SendKeys{target: format!("{}:{}", project_name, k.as_str().unwrap().to_string()).to_string(), exec: v.as_str().expect("Bad exec command").to_string()}));
                        }
                    }
                },
                &Yaml::String(ref s) => {
                    commands.push(Command::Window(Window{session_name: format!("{}:{}", project_name, i+1), name: s.clone(), root: root.clone()}))
                },
                &Yaml::Integer(ref s) => {
                    commands.push(Command::Window(Window{session_name: format!("{}:{}", project_name, i+1), name: s.to_string(), root: root.clone()}))
                },
                _ => panic!("Muxed config file formatting isn't recognized.")
            };
        };
    };

    commands.push(Command::KillWindow(KillWindow{name: format!("{}:{}", project_name, tmp_window_name)}));

    if !daemonize { commands.push(Command::Attach(Attach{name: project_name.clone()})) };
    Ok(commands)
}

/// Pane matcher is for breaking apart the panes. Splitting windows when needed
/// and executing commands as needed.
fn pane_matcher(session: &String, window: &Yaml, root: &Option<String>, window_name: String) -> Result<Vec<Command>, String> {
    let mut commands = vec!();
    let panes = window["panes"].as_vec().expect("Something is wrong with panes.");

    for (i, pane) in panes.iter().enumerate() {
        // For every pane, we need one less split.
        // ex. An existing window to become 2 panes, needs 1 split.
        if i < (panes.len()-1) {
            commands.push(Command::Split(Split{target: format!("{}:{}.{}", session, window_name, i).to_string(), root: root.clone()}));
        };
        // Execute given commands in each new pane after all splits are
        // complete.
        commands.push(Command::SendKeys(SendKeys{target: format!("{}:{}.{}", session, window_name, i).to_string(), exec: pane.as_str().expect("Bad exec command").to_string()}));
    };

    // After all panes are split select the layout for the window
    if window["layout"].as_str().is_some() {
        let layout = window["layout"].as_str().expect("Bad layout").to_string();
        commands.push(Command::Layout(Layout{target: format!("{}:{}", session, window_name).to_string(), layout: layout}));
    };

    Ok(commands)
}

#[test]
pub fn windows_defined_as_array_has_6_commands() {
    let s = "---
windows: ['cargo', 'vim', 'git']
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    assert_eq!(main(&yaml, &"muxed".to_string(), false).unwrap().len(), 6)
}

#[test]
pub fn windows_defined_as_array_has_1_session() {
    let s = "---
windows: ['cargo', 'vim', 'git']
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string(), false).unwrap().into_iter().filter(|x| match x {
        &Command::Session(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 1)
}

#[test]
pub fn windows_defined_as_array_has_3_windows() {
    let s = "---
windows: ['cargo', 'vim', 'git']
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string(), false).unwrap().into_iter().filter(|x| match x {
        &Command::Window(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 3)
}

#[test]
pub fn windows_defined_as_array_has_1_attach() {
    let s = "---
windows: ['cargo', 'vim', 'git']
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string(), false).unwrap().into_iter().filter(|x| match x {
        &Command::Attach(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 1)
}

#[test]
pub fn windows_with_integer_names() {
    let s = "---
windows: [1, 'vim', 3]
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    assert_eq!(main(&yaml, &"muxed".to_string(), false).unwrap().len(), 6)
}

#[test]
pub fn windows_as_list() {
    let s = "---
windows:
  - cargo: ''
  - vim: ''
  - git: ''
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let commands = main(&yaml, &"muxed".to_string(), false).unwrap();
    assert_eq!(commands.len(), 9)
}

#[test]
pub fn panes_array_has_8_commands() {
    let s = "---
windows:
  - editor:
      layout: 'main-vertical'
      panes: ['vim', 'guard']
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let commands = main(&yaml, &"muxed".to_string(), false).unwrap();
    assert_eq!(commands.len(), 8)
}

#[test]
pub fn panes_array_has_1_split() {
    let s = "---
windows:
  - editor:
      layout: 'main-vertical'
      panes: ['vim', 'guard']
";

    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string(), false).unwrap().into_iter().filter(|x| match x {
        &Command::Split(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 1)
}

#[test]
pub fn panes_array_has_1_layout() {
    let s = "---
windows:
  - editor:
      layout: 'main-vertical'
      panes: ['vim', 'guard']
";

    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string(), false).unwrap().into_iter().filter(|x| match x {
        &Command::Layout(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 1)
}

#[test]
pub fn panes_array_has_1_window() {
    let s = "---
windows:
  - editor:
      layout: 'main-vertical'
      panes: ['vim', 'guard']
";

    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string(), false).unwrap().into_iter().filter(|x| match x {
        &Command::Window(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 1)
}

#[test]
pub fn panes_array_has_1_session() {
    let s = "---
windows:
  - editor:
      layout: 'main-vertical'
      panes: ['vim', 'guard']
";

    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string(), false).unwrap().into_iter().filter(|x| match x {
        &Command::Session(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 1)
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::fmt::Display;
use std::fs::{File, OpenOptions};
use std::io::prelude::*;
use std::io::BufWriter;
use std::mem;
use std::net::{UdpSocket, IpAddr};
use std::time::{Duration, Instant};

extern crate flate2;
use flate2::write::ZlibEncoder;
use flate2::Compression;

// This server indiscriminately saves any packet to a compressed, rotated,
// per-exporter file. No checking or parsing of PDUs is performed.

fn main() {
    let compression = Compression::Best;
    let rotate_interval = Duration::from_secs(300);
    let mut file_map: HashMap<IpAddr, FlowFileWriter<IpAddr, _>> = HashMap::new();
    let socket = UdpSocket::bind("[::]:4013").unwrap();
    let mut recv_buf = [0u8; 2048];

    while let Ok((amt, src)) = socket.recv_from(&mut recv_buf[..]) {
        let mut output = file_map.entry(src.ip()).or_insert_with(|| {
            println!("New exporter: {}", src.ip());
            FlowFileWriter::compressed(src.ip(), compression, rotate_interval).unwrap()
        });

        output.maybe_rotate().unwrap();
        output.write(&recv_buf[..amt]).unwrap();
    }
}

#[derive(Debug)]
enum Error {
    Io(std::io::Error)
}

struct FlowFileWriter<D, W> where D: Display + Copy, W: Write {
    exporter_name: D,
    generation: usize,
    compression: Compression,
    last_rotation: Instant,
    rotation_interval: Duration,
    output: W,
}

type Output = BufWriter<ZlibEncoder<File>>;

impl<D> FlowFileWriter<D, Output> where D: Display + Copy {

    fn compressed(exporter_name: D, compression: Compression, interval: Duration)
           -> Result<FlowFileWriter<D, Output>, Error> {
        Ok(
            FlowFileWriter {
                exporter_name: exporter_name,
                generation: 0,
                compression: compression,
                last_rotation: Instant::now(),
                rotation_interval: interval,
                output: BufWriter::new(
                    ZlibEncoder::new(
                        try!(OpenOptions::new()
                             .create(true)
                             .append(true)
                             .open(format!("./{}.{:04}.z", exporter_name, 0))
                             .map_err(|err| { Error::Io(err) })), compression))
            }
        )
    }

    fn rotate(&mut self) -> Result<(), Error> {
        self.generation += 1;
        self.output = BufWriter::new(ZlibEncoder::new(
            try!(OpenOptions::new()
                 .create(true)
                 .append(true)
                 .open(format!("./{}.{:04}.z", self.exporter_name, self.generation))
                 .map_err(|err| { Error::Io(err) })), self.compression));
        self.last_rotation = Instant::now();
        Ok(())
    }

    fn maybe_rotate(&mut self) -> Result<(), Error> {
        if self.last_rotation.elapsed() > self.rotation_interval {
            self.rotate()
        } else {
            Ok(())
        }
    }
}

impl<D,W> Write for FlowFileWriter<D, W> where D: Display + Copy, W: Write {

    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        let length: [u8; 2] = unsafe { mem::transmute((buf.len() as u16).to_be()) };
        try!(self.output.write(&length[..]));
        self.output.write(buf)
    }

    fn flush(&mut self) -> std::io::Result<()>{
        self.output.flush()
    }
}
<commit_msg>New file format for "record".<commit_after>use std::collections::{HashMap, HashSet};
use std::fmt::Display;
use std::fs::{File, OpenOptions};
use std::io::prelude::*;
use std::io::{BufWriter, SeekFrom};
use std::mem;
use std::net::{UdpSocket, IpAddr};
use std::time::{Duration, Instant};

extern crate flate2;
use flate2::write::ZlibEncoder;
use flate2::Compression;

extern crate netflowv9;
use netflowv9::*;

fn main() {
    let compression = Compression::Best;
    let rotate_interval = Duration::from_secs(30);
    let mut file_map: HashMap<IpAddr, FlowFileWriter<IpAddr, _>> = HashMap::new();
    let socket = UdpSocket::bind("[::]:4013").unwrap();
    let mut recv_buf = [0u8; 2048];

    while let Ok((amt, src)) = socket.recv_from(&mut recv_buf[..]) {
        let mut output = file_map.entry(src.ip()).or_insert_with(|| {
            println!("New exporter: {}", src.ip());
            FlowFileWriter::compressed(src.ip(), compression, rotate_interval).unwrap()
        });

        output.maybe_rotate().unwrap();
        output.write(&recv_buf[..amt]).unwrap();
    }
}

type RawDataTemplateIter = Vec<u8>;

struct FlowFileWriter<D, W> where D: Display + Copy, W: Write {
    exporter_name: D,
    generation: usize,
    compression: Compression,
    last_rotation: Instant,
    rotation_interval: Duration,
    template_ids_seen: HashSet<u16>,
    templates: HashMap<u16, Vec<u8>>,
    output: W,
}

type Output = BufWriter<ZlibEncoder<File>>;

impl<D> FlowFileWriter<D, Output> where D: Display + Copy {

    fn compressed(exporter_name: D, compression: Compression, interval: Duration)
           -> Result<FlowFileWriter<D, Output>, std::io::Error> {
        let file = try!(Self::init_file(exporter_name, 0));
        Ok(
            FlowFileWriter {
                exporter_name: exporter_name,
                generation: 0,
                compression: compression,
                last_rotation: Instant::now(),
                rotation_interval: interval,
                template_ids_seen: HashSet::new(),
                templates: HashMap::new(),
                output: BufWriter::new(ZlibEncoder::new(file, compression))
            }
        )
    }

    fn init_file(exporter_name: D, generation: usize) -> Result<File, std::io::Error> {
        let mut file = try!(OpenOptions::new().create(true).write(true).read(true)
                            .open(format!("./{}.{:04}", exporter_name, generation)));
        // magic
        try!(file.write(b"NETFLO00"));
        // offset to metadata packets (updated in finalize_file)
        try!(file.write(&[0; 8]));
        Ok(file)
    }

    /// Dump templates needed to decode flow records at end of file.
    /// Update template offset in header (bytes 8 - 15).
    fn finalize_file(&mut self, mut f: File) -> Result<(), std::io::Error> {
        let file_len = f.metadata().unwrap().len();
        let len: [u8; 8] = unsafe { mem::transmute((file_len - 16).to_be())};
        assert!(file_len > 0);

        f.seek(SeekFrom::Start(8)).unwrap();
        f.write(&len[..]).unwrap();
        f.seek(SeekFrom::Start(file_len)).unwrap();

        for template_id in &self.template_ids_seen {
            if let Some(raw_template) = self.templates.get(&template_id) {
                println!("saving template {}", template_id);
                unsafe {
                    let tid: [u8;2] = mem::transmute(template_id.to_be());
                    f.write(&tid[..]).unwrap();
                    let len: [u8;2] = mem::transmute((raw_template.len() as u16).to_be());
                    f.write(&len[..]).unwrap();
                }
                f.write(raw_template).unwrap();
            } else {
                println!("missing template {}", template_id);
            }
        }

        self.template_ids_seen.clear();

        Ok(())
    }

    fn rotate(&mut self) -> Result<(), std::io::Error> {
        println!("Rotating {}.{}", self.exporter_name, self.generation);
        self.generation += 1;
        self.last_rotation = Instant::now();

        let new_writer =  BufWriter::new(ZlibEncoder::new(
            try!(Self::init_file(self.exporter_name, self.generation)), self.compression));
        let writer = mem::replace(&mut self.output, new_writer);
        let zlibencoder = writer.into_inner().ok().unwrap();

        let old_file = zlibencoder.finish().unwrap();
        self.finalize_file(old_file)
    }

    fn maybe_rotate(&mut self) -> Result<(), std::io::Error> {
        if self.last_rotation.elapsed() > self.rotation_interval {
            self.rotate()
        } else {
            Ok(())
        }
    }
}

impl<D,W> Write for FlowFileWriter<D, W> where D: Display + Copy, W: Write {

    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        let packet = match Packet::from_bytes(buf) {
            Ok(packet) => packet,
            Err(_) => return Ok(buf.len()),
        };

        for flowset in packet.flowsets() {
            match flowset {
                FlowSet::OptionsTemplates(templates) => {
                    for template in templates.clone() {
                        if self.templates.get(&template.template_id).is_none() {
                            self.templates.insert(template.template_id,
                                                  template.fields.raw.to_vec());
                        }
                    }
                }
                FlowSet::DataTemplates(templates) => {
                    for template in templates.clone() {
                        if self.templates.get(&template.template_id).is_none() {
                            self.templates.insert(template.template_id,
                                                  template.fields.raw.to_vec());
                        }
                    }
                }
                FlowSet::DataRecords(records) => {
                    self.template_ids_seen.insert(records.template_id);
                    let template_id: [u8; 2] = unsafe {
                        mem::transmute((records.template_id).to_be())
                    };
                    try!(self.output.write(&template_id[..]));
                    let length: [u8; 2] = unsafe {
                        mem::transmute((buf.len() as u16).to_be())
                    };
                    try!(self.output.write(&length[..]));
                    try!(self.output.write(records.raw));
                }
                _ => ()
            }
        }

        Ok(buf.len())
    }

    fn flush(&mut self) -> std::io::Result<()>{
        self.output.flush()
    }
}
<|endoftext|>


<commit_before>use luminance::tessellation;
use luminance_gl::gl33::Tessellation;
use std::collections::BTreeMap;
use std::fs::File;
use std::io::Read;
use std::iter::IntoIterator;
use std::path::{Path, PathBuf};
use std::vec;
use wavefront_obj::{self, obj};

use resource::{Cache, Load, LoadError};

pub type Vertex = (VertexPos, VertexNor, VertexTexCoord);
pub type VertexPos = [f32; 3];
pub type VertexNor = [f32; 3];
pub type VertexTexCoord = [f32; 2];

pub struct Model {
  pub parts: Vec<Part>
}

impl Model {
  pub fn from_parts(parts: Vec<Part>) -> Self {
    Model {
      parts: parts
    }
  }
}

impl IntoIterator for Model {
  type Item = Part;
  type IntoIter = vec::IntoIter<Part>;

  fn into_iter(self) -> Self::IntoIter {
    self.parts.into_iter()
  }
}

pub struct Part {
  pub tess: Tessellation,
  // TODO: add material index
}

impl Part {
  pub fn new(tess: Tessellation) -> Self {
    Part {
      tess: tess,
    }
  }
}

impl<'a> Load<'a> for Model {
  type Args = ();

  fn load<P>(path: P, _: &mut Cache<'a>, _: Self::Args) -> Result<Self, LoadError> where P: AsRef<Path> {
    let path = path.as_ref();

    info!("loading model: {:?}", path);

    let mut input = String::new();

    // load the data directly into memory; no buffering nor streaming
    {
      let mut file = File::open(path).map_err(|e| LoadError::FileNotFound(path.to_path_buf(), format!("{:?}", e)))?;
      let _ = file.read_to_string(&mut input);
    }

    // parse the obj file and convert it
    let obj_set = obj::parse(input).map_err(|e| LoadError::ParseFailed(format!("{:?}", e)))?;

    convert_obj(obj_set).map_err(|e| LoadError::ConversionFailed(format!("{:?}", e)))
  }
}

// Turn a wavefront obj object into a `Model`
fn convert_obj(obj_set: obj::ObjSet) -> Result<Model, ModelError> {
  let mut parts = Vec::new();

  info!("{} objects to convert…", obj_set.objects.len());
  for obj in &obj_set.objects {
    info!("  converting {} geometries in object {}", obj.geometry.len(), obj.name);

    // convert all the geometries
    for geometry in &obj.geometry {
      info!("    {} vertices, {} normals, {} tex vertices", obj.vertices.len(), obj.normals.len(), obj.tex_vertices.len());
      let (vertices, indices, mode) = convert_geometry(geometry, &obj.vertices, &obj.normals, &obj.tex_vertices)?;
      let part = Part::new(Tessellation::new(mode, &vertices, Some(&indices))); // FIXME: material
      parts.push(part);
    }
  }

  Ok(Model::from_parts(parts))
}

// Convert wavefront_obj’s Geometry into a pair of vertices and indices.
//
// This function will regenerate the indices on the fly based on which are used in the shapes in the
// geometry. It’s used to create independent tessellation.
fn convert_geometry(geo: &obj::Geometry, positions: &[obj::Vertex], normals: &[obj::Normal], tvertices: &[obj::TVertex]) -> Result<(Vec<Vertex>, Vec<u32>, tessellation::Mode), ModelError> {
  if geo.shapes.is_empty() {
    return Err(ModelError::NoShape);
  }

  let mut vertices = Vec::new(); // FIXME: better allocation scheme?
  let mut indices = Vec::new();
  let mut index_map = BTreeMap::new();

  info!("    converting geometry");

  let mode = guess_mode(geo.shapes[0].primitive);

  for prim in geo.shapes.iter().map(|s| s.primitive) {
    let keys = create_keys_from_primitive(prim)?;

    for key in keys {
      match index_map.get(&key).map(|&i| i) {
        Some(index) => {
          // that triplet already exists; just append the index in the indices buffer
          indices.push(index);
        },
        None => {
          // this is a new, not yet discovered triplet; create the corresponding vertex and add it
          // to the vertices buffer, and map the triplet to the index in the indices buffer
          let vertex = interleave_vertex(&positions[key.0], &normals[key.1], key.2.map(|ki| &tvertices[ki]));
          let index = vertices.len() as u32;

          vertices.push(vertex);
          indices.push(index);
          index_map.insert(key, index);
        }
      }
    }
  }

  Ok((vertices, indices, mode))
}

// Create triplet keys from wavefront_obj primitives. If any primitive doesn’t have all the triplet
// information (position, normal, tex), a ModelError::UnsupportedVertex error is returned instead.
fn create_keys_from_primitive(prim: obj::Primitive) -> Result<Vec<(usize, usize, Option<usize>)>, ModelError> {
  match prim {
    obj::Primitive::Point(i) => {
      let a = vtnindex_to_key(i)?;
      Ok(vec![a])
    },
    obj::Primitive::Line(i, j) => {
      let a = vtnindex_to_key(i)?;
      let b = vtnindex_to_key(j)?;
      Ok(vec![a, b])
    },
    obj::Primitive::Triangle(i, j, k) => {
      let a = vtnindex_to_key(i)?;
      let b = vtnindex_to_key(j)?;
      let c = vtnindex_to_key(k)?;
      Ok(vec![a, b, c])
    }
  }
}

// Convert from a wavefront_obj VTNIndex into our triplet, raising error if not possible.
fn vtnindex_to_key(i: obj::VTNIndex) -> Result<(usize, usize, Option<usize>), ModelError> {
  match i {
    (pi, ti, Some(ni)) => Ok((pi, ni, ti)),
    _ => Err(ModelError::UnsupportedVertex)
  }
}

fn interleave_vertex(p: &obj::Vertex, n: &obj::Normal, t: Option<&obj::TVertex>) -> Vertex {
  (convert_vertex(p), convert_nor(n), t.map_or([0., 0.], convert_tvertex))
}

fn convert_vertex(v: &obj::Vertex) -> VertexPos {
  [v.x as f32, v.y as f32, v.z as f32]
}

fn convert_nor(n: &obj::Normal) -> VertexNor {
  convert_vertex(n)
}

fn convert_tvertex(t: &obj::TVertex) -> VertexTexCoord {
  [t.u as f32, t.v as f32]
}

fn guess_mode(prim: obj::Primitive) -> tessellation::Mode {
  match prim {
    obj::Primitive::Point(_) => tessellation::Mode::Point,
    obj::Primitive::Line(_, _) => tessellation::Mode::Line,
    obj::Primitive::Triangle(_, _, _) => tessellation::Mode::Triangle
  }
}

#[derive(Debug)]
pub enum ModelError {
  FileNotFound(PathBuf, String),
  ParseFailed(wavefront_obj::ParseError),
  MultiObjects,
  UnsupportedVertex,
  NoShape
}
<commit_msg>Removed useless warnings.<commit_after>use luminance::tessellation;
use luminance_gl::gl33::Tessellation;
use std::collections::BTreeMap;
use std::fs::File;
use std::io::Read;
use std::iter::IntoIterator;
use std::path::Path;
use std::vec;
use wavefront_obj::obj;

use resource::{Cache, Load, LoadError};

pub type Vertex = (VertexPos, VertexNor, VertexTexCoord);
pub type VertexPos = [f32; 3];
pub type VertexNor = [f32; 3];
pub type VertexTexCoord = [f32; 2];

pub struct Model {
  pub parts: Vec<Part>
}

impl Model {
  pub fn from_parts(parts: Vec<Part>) -> Self {
    Model {
      parts: parts
    }
  }
}

impl IntoIterator for Model {
  type Item = Part;
  type IntoIter = vec::IntoIter<Part>;

  fn into_iter(self) -> Self::IntoIter {
    self.parts.into_iter()
  }
}

pub struct Part {
  pub tess: Tessellation,
  // TODO: add material index
}

impl Part {
  pub fn new(tess: Tessellation) -> Self {
    Part {
      tess: tess,
    }
  }
}

impl<'a> Load<'a> for Model {
  type Args = ();

  fn load<P>(path: P, _: &mut Cache<'a>, _: Self::Args) -> Result<Self, LoadError> where P: AsRef<Path> {
    let path = path.as_ref();

    info!("loading model: {:?}", path);

    let mut input = String::new();

    // load the data directly into memory; no buffering nor streaming
    {
      let mut file = File::open(path).map_err(|e| LoadError::FileNotFound(path.to_path_buf(), format!("{:?}", e)))?;
      let _ = file.read_to_string(&mut input);
    }

    // parse the obj file and convert it
    let obj_set = obj::parse(input).map_err(|e| LoadError::ParseFailed(format!("{:?}", e)))?;

    convert_obj(obj_set).map_err(|e| LoadError::ConversionFailed(format!("{:?}", e)))
  }
}

// Turn a wavefront obj object into a `Model`
fn convert_obj(obj_set: obj::ObjSet) -> Result<Model, ModelError> {
  let mut parts = Vec::new();

  info!("{} objects to convert…", obj_set.objects.len());
  for obj in &obj_set.objects {
    info!("  converting {} geometries in object {}", obj.geometry.len(), obj.name);

    // convert all the geometries
    for geometry in &obj.geometry {
      info!("    {} vertices, {} normals, {} tex vertices", obj.vertices.len(), obj.normals.len(), obj.tex_vertices.len());
      let (vertices, indices, mode) = convert_geometry(geometry, &obj.vertices, &obj.normals, &obj.tex_vertices)?;
      let part = Part::new(Tessellation::new(mode, &vertices, Some(&indices))); // FIXME: material
      parts.push(part);
    }
  }

  Ok(Model::from_parts(parts))
}

// Convert wavefront_obj’s Geometry into a pair of vertices and indices.
//
// This function will regenerate the indices on the fly based on which are used in the shapes in the
// geometry. It’s used to create independent tessellation.
fn convert_geometry(geo: &obj::Geometry, positions: &[obj::Vertex], normals: &[obj::Normal], tvertices: &[obj::TVertex]) -> Result<(Vec<Vertex>, Vec<u32>, tessellation::Mode), ModelError> {
  if geo.shapes.is_empty() {
    return Err(ModelError::NoShape);
  }

  let mut vertices = Vec::new(); // FIXME: better allocation scheme?
  let mut indices = Vec::new();
  let mut index_map = BTreeMap::new();

  info!("    converting geometry");

  let mode = guess_mode(geo.shapes[0].primitive);

  for prim in geo.shapes.iter().map(|s| s.primitive) {
    let keys = create_keys_from_primitive(prim)?;

    for key in keys {
      match index_map.get(&key).map(|&i| i) {
        Some(index) => {
          // that triplet already exists; just append the index in the indices buffer
          indices.push(index);
        },
        None => {
          // this is a new, not yet discovered triplet; create the corresponding vertex and add it
          // to the vertices buffer, and map the triplet to the index in the indices buffer
          let vertex = interleave_vertex(&positions[key.0], &normals[key.1], key.2.map(|ki| &tvertices[ki]));
          let index = vertices.len() as u32;

          vertices.push(vertex);
          indices.push(index);
          index_map.insert(key, index);
        }
      }
    }
  }

  Ok((vertices, indices, mode))
}

// Create triplet keys from wavefront_obj primitives. If any primitive doesn’t have all the triplet
// information (position, normal, tex), a ModelError::UnsupportedVertex error is returned instead.
fn create_keys_from_primitive(prim: obj::Primitive) -> Result<Vec<(usize, usize, Option<usize>)>, ModelError> {
  match prim {
    obj::Primitive::Point(i) => {
      let a = vtnindex_to_key(i)?;
      Ok(vec![a])
    },
    obj::Primitive::Line(i, j) => {
      let a = vtnindex_to_key(i)?;
      let b = vtnindex_to_key(j)?;
      Ok(vec![a, b])
    },
    obj::Primitive::Triangle(i, j, k) => {
      let a = vtnindex_to_key(i)?;
      let b = vtnindex_to_key(j)?;
      let c = vtnindex_to_key(k)?;
      Ok(vec![a, b, c])
    }
  }
}

// Convert from a wavefront_obj VTNIndex into our triplet, raising error if not possible.
fn vtnindex_to_key(i: obj::VTNIndex) -> Result<(usize, usize, Option<usize>), ModelError> {
  match i {
    (pi, ti, Some(ni)) => Ok((pi, ni, ti)),
    _ => Err(ModelError::UnsupportedVertex)
  }
}

fn interleave_vertex(p: &obj::Vertex, n: &obj::Normal, t: Option<&obj::TVertex>) -> Vertex {
  (convert_vertex(p), convert_nor(n), t.map_or([0., 0.], convert_tvertex))
}

fn convert_vertex(v: &obj::Vertex) -> VertexPos {
  [v.x as f32, v.y as f32, v.z as f32]
}

fn convert_nor(n: &obj::Normal) -> VertexNor {
  convert_vertex(n)
}

fn convert_tvertex(t: &obj::TVertex) -> VertexTexCoord {
  [t.u as f32, t.v as f32]
}

fn guess_mode(prim: obj::Primitive) -> tessellation::Mode {
  match prim {
    obj::Primitive::Point(_) => tessellation::Mode::Point,
    obj::Primitive::Line(_, _) => tessellation::Mode::Line,
    obj::Primitive::Triangle(_, _, _) => tessellation::Mode::Triangle
  }
}

#[derive(Debug)]
pub enum ModelError {
  UnsupportedVertex,
  NoShape
}
<|endoftext|>