<commit_before>/// Implement +, -, *, / operators for Tensor
/// +=, -= are not supported but do as ops::inplace_add, ops::inplace_sub
extern crate ndarray;

use ndarray_ext::NdArray;
use ops;
use std::mem;
use std::ops::{Add, Div, Mul, Sub};
use std::rc::Rc;
use tensor::{RawTensor, Tensor};

#[inline(always)]
fn scalar_tensor_add(a: &NdArray, b: &NdArray) -> NdArray
{
    // comparing the rank of tensors doesn't solve the problem here.
    if a.len() > b.len() { a + b } else { b + a }
}

#[inline(always)]
fn scalar_tensor_mul(a: &NdArray, b: &NdArray) -> NdArray
{
    // comparing the rank of tensors doesn't solve the problem here.
    if a.len() > b.len() { a * b } else { b * a }
}

pub struct AddOp;
pub struct SubOp;
pub struct MulOp;
pub struct DivOp;

pub struct InplaceAddOp;
pub struct InplaceSubOp;

impl ops::Op for AddOp {
    fn name(&self) -> &str
    {
        "AddOp"
    }

    fn compute(&self, xs: &[&NdArray], _: bool) -> NdArray
    {
        let a = xs[0];
        let b = xs[1];
        scalar_tensor_add(a, b)
    }

    fn grad(&self, gy: &Tensor, _: &[&Tensor], _: &Tensor) -> Vec<Option<Tensor>>
    {
        vec![Some(gy.clone()), Some(gy.clone())]
    }
}

impl ops::Op for SubOp {
    fn name(&self) -> &str
    {
        "SubOp"
    }

    fn compute(&self, xs: &[&NdArray], _: bool) -> NdArray
    {
        let a = xs[0];
        let b = xs[1];
        if a.shape() == &[1] {
            // a is scalar
            // unwrap is safe
            let a = NdArray::from_elem(b.shape(), *a.get(0).unwrap());
            a - b
        } else {
            a - b
        }
    }

    fn grad(&self, gy: &Tensor, _: &[&Tensor], _: &Tensor) -> Vec<Option<Tensor>>
    {
        vec![Some(gy.clone()), Some(-1 * gy)]
    }
}

impl ops::Op for MulOp {
    fn name(&self) -> &str
    {
        "MulOp"
    }

    fn compute(&self, xs: &[&NdArray], _: bool) -> NdArray
    {
        let a = xs[0];
        let b = xs[1];
        scalar_tensor_mul(a, b)
    }

    fn grad(&self, gy: &Tensor, inputs: &[&Tensor], _: &Tensor) -> Vec<Option<Tensor>>
    {
        let x0 = inputs[0];
        let x1 = inputs[1];
        vec![Some(gy * x1), Some(gy * x0)]
    }
}

impl ops::Op for DivOp {
    fn name(&self) -> &str
    {
        "DivOp"
    }

    fn compute(&self, xs: &[&NdArray], _: bool) -> NdArray
    {
        let x0 = xs[0];
        let x1 = xs[1];
        if x0.shape() == &[1] {
            // a is scalar
            // unwrap is safe
            NdArray::from_elem(x1.shape(), *x0.get(0).unwrap()) / x1
        } else {
            x0 / x1
        }
    }

    fn grad(&self, gy: &Tensor, inputs: &[&Tensor], _: &Tensor) -> Vec<Option<Tensor>>
    {
        let x0 = inputs[0];
        let x1 = inputs[1];
        let a: Tensor = gy / x1;
        let b: Tensor = -1 * x0 * ops::pow(x1, -2.) * gy;
        vec![Some(a), Some(b)]
    }
}

impl ops::Op for InplaceAddOp {
    fn inplace(&self) -> bool
    {
        true
    }

    fn name(&self) -> &str
    {
        "InplaceAddOp"
    }

    fn compute_inplace(&self, xs: &mut [&mut NdArray], _: bool)
    {
        // safe transmute
        let b: &&NdArray = unsafe { mem::transmute(&mut xs[1]) };
        let a = &mut xs[0];
        a.zip_mut_with(b, |x, &y| *x += y);
    }

    fn grad(&self, gy: &Tensor, _: &[&Tensor], _: &Tensor) -> Vec<Option<Tensor>>
    {
        vec![Some(gy.clone()), Some(gy.clone())]
    }
}

impl ops::Op for InplaceSubOp {
    fn inplace(&self) -> bool
    {
        true
    }

    fn name(&self) -> &str
    {
        "InplaceSubOp"
    }

    fn compute_inplace(&self, xs: &mut [&mut NdArray], _: bool)
    {
        // safe transmute
        let b: &&NdArray = unsafe { mem::transmute(&mut xs[1]) };
        let a = &mut xs[0];
        a.zip_mut_with(b, |x, &y| *x -= y);
    }

    fn grad(&self, gy: &Tensor, _: &[&Tensor], _: &Tensor) -> Vec<Option<Tensor>>
    {
        vec![Some(gy.clone()), Some(-1 * gy)]
    }
}

#[inline]
fn scalar_to_tensor(arg: f32) -> Tensor
{
    Tensor(Rc::new(RawTensor {
        op: Box::new(ops::scalar::Scalar { val: arg }),
        inputs: vec![],
        top_rank: 0,
    }))
}

#[inline]
fn f32_to_tensor(arg: f32) -> Tensor
{
    scalar_to_tensor(arg)
}

#[inline]
fn f64_to_tensor(arg: f64) -> Tensor
{
    scalar_to_tensor(arg as f32)
}

#[inline]
fn i32_to_tensor(arg: i32) -> Tensor
{
    scalar_to_tensor(arg as f32)
}

#[inline]
fn i64_to_tensor(arg: i64) -> Tensor
{
    scalar_to_tensor(arg as f32)
}

#[inline]
fn u32_to_tensor(arg: u32) -> Tensor
{
    scalar_to_tensor(arg as f32)
}

#[inline]
fn u64_to_tensor(arg: u64) -> Tensor
{
    scalar_to_tensor(arg as f32)
}

#[inline]
fn usize_to_tensor(arg: usize) -> Tensor
{
    scalar_to_tensor(arg as f32)
}

#[inline]
fn isize_to_tensor(arg: isize) -> Tensor
{
    scalar_to_tensor(arg as f32)
}

// -- std::ops::{Add, Sub, Mul, Div} implementations --

macro_rules! impl_elementwise_between_tensor_and_scalar {
    (
        $trt:ident,
        $func:ident,
        $op:ident,
        $scalar_type:ty,
        $conversion_fn:ident
    ) => {

        // scalar op Tensor
        impl $trt<Tensor> for $scalar_type {
            type Output = Tensor;
            fn $func(self, rhs: Tensor) -> Self::Output
            {
                ops::apply_op($op, &[&$conversion_fn(self), &rhs])
            }
        }

        // scalar op &Tensor
        impl<'a> $trt<&'a Tensor> for $scalar_type {
            type Output = Tensor;
            fn $func(self, rhs: &'a Tensor) -> Self::Output
            {
                ops::apply_op($op, &[&$conversion_fn(self), rhs])
            }
        }

        // Tensor op scalar
        impl $trt<$scalar_type> for Tensor {
            type Output = Tensor;
            fn $func(self, rhs: $scalar_type) -> Self::Output
            {
                ops::apply_op($op, &[&self, &$conversion_fn(rhs)])
            }
        }

        // &Tensor op scalar
        impl<'a> $trt<$scalar_type> for &'a Tensor {
            type Output = Tensor;
            fn $func(self, rhs: $scalar_type) -> Self::Output
            {
                ops::apply_op($op, &[self, &$conversion_fn(rhs)])
            }
        }
    }
}

macro_rules! impl_elementwise_between_two_tensors {
    (
        $trt:ident,
        $func:ident,
        $op:ident
    ) => {
        // Tensor op Tensor
        impl $trt for Tensor {
            type Output = Tensor;
            fn $func(self, rhs: Tensor) -> Self::Output
            {
                ops::apply_op($op, &[&self, &rhs])
            }
        }

        // Tensor op &Tensor
        impl<'a> $trt<&'a Tensor> for Tensor {
            type Output = Tensor;
            fn $func(self, rhs: &Tensor) -> Self::Output
            {
                ops::apply_op($op, &[&self, rhs])
            }
        }

        // &Tensor op Tensor
        impl<'a> $trt<Tensor> for &'a Tensor {
            type Output = Tensor;
            fn $func(self, rhs: Tensor) -> Self::Output
            {
                ops::apply_op($op, &[self, &rhs])
            }
        }

        // &Tensor op &Tensor
        // lifetime of the two tensors are unrelated
        impl<'a, 'b> $trt<&'a Tensor> for &'b Tensor {
            type Output = Tensor;
            fn $func(self, rhs: &Tensor) -> Self::Output
            {
                ops::apply_op($op, &[self, rhs])
            }
        }
    };
}

impl_elementwise_between_two_tensors!(Add, add, AddOp);
impl_elementwise_between_two_tensors!(Sub, sub, SubOp);
impl_elementwise_between_two_tensors!(Mul, mul, MulOp);
impl_elementwise_between_two_tensors!(Div, div, DivOp);

impl_elementwise_between_tensor_and_scalar!(Add, add, AddOp, i32, i32_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Sub, sub, SubOp, i32, i32_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Mul, mul, MulOp, i32, i32_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Div, div, DivOp, i32, i32_to_tensor);

impl_elementwise_between_tensor_and_scalar!(Add, add, AddOp, i64, i64_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Sub, sub, SubOp, i64, i64_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Mul, mul, MulOp, i64, i64_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Div, div, DivOp, i64, i64_to_tensor);

impl_elementwise_between_tensor_and_scalar!(Add, add, AddOp, f32, f32_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Sub, sub, SubOp, f32, f32_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Mul, mul, MulOp, f32, f32_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Div, div, DivOp, f32, f32_to_tensor);

impl_elementwise_between_tensor_and_scalar!(Add, add, AddOp, f64, f64_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Sub, sub, SubOp, f64, f64_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Mul, mul, MulOp, f64, f64_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Div, div, DivOp, f64, f64_to_tensor);

impl_elementwise_between_tensor_and_scalar!(Add, add, AddOp, u32, u32_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Sub, sub, SubOp, u32, u32_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Mul, mul, MulOp, u32, u32_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Div, div, DivOp, u32, u32_to_tensor);

impl_elementwise_between_tensor_and_scalar!(Add, add, AddOp, u64, u64_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Sub, sub, SubOp, u64, u64_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Mul, mul, MulOp, u64, u64_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Div, div, DivOp, u64, u64_to_tensor);

impl_elementwise_between_tensor_and_scalar!(Add, add, AddOp, usize, usize_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Sub, sub, SubOp, usize, usize_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Mul, mul, MulOp, usize, usize_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Div, div, DivOp, usize, usize_to_tensor);

impl_elementwise_between_tensor_and_scalar!(Add, add, AddOp, isize, isize_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Sub, sub, SubOp, isize, isize_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Mul, mul, MulOp, isize, isize_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Div, div, DivOp, isize, isize_to_tensor);
<commit_msg>Use macro for scalar_to_tensor<commit_after>/// Implement +, -, *, / operators for Tensor
/// +=, -= are not supported but do as ops::inplace_add, ops::inplace_sub
extern crate ndarray;

use ndarray_ext::NdArray;
use ops;
use std::mem;
use std::ops::{Add, Div, Mul, Sub};
use std::rc::Rc;
use tensor::{RawTensor, Tensor};

#[inline(always)]
fn scalar_tensor_add(a: &NdArray, b: &NdArray) -> NdArray
{
    // comparing the rank of tensors doesn't solve the problem here.
    if a.len() > b.len() { a + b } else { b + a }
}

#[inline(always)]
fn scalar_tensor_mul(a: &NdArray, b: &NdArray) -> NdArray
{
    // comparing the rank of tensors doesn't solve the problem here.
    if a.len() > b.len() { a * b } else { b * a }
}

pub struct AddOp;
pub struct SubOp;
pub struct MulOp;
pub struct DivOp;

pub struct InplaceAddOp;
pub struct InplaceSubOp;

impl ops::Op for AddOp {
    fn name(&self) -> &str
    {
        "AddOp"
    }

    fn compute(&self, xs: &[&NdArray], _: bool) -> NdArray
    {
        let a = xs[0];
        let b = xs[1];
        scalar_tensor_add(a, b)
    }

    fn grad(&self, gy: &Tensor, _: &[&Tensor], _: &Tensor) -> Vec<Option<Tensor>>
    {
        vec![Some(gy.clone()), Some(gy.clone())]
    }
}

impl ops::Op for SubOp {
    fn name(&self) -> &str
    {
        "SubOp"
    }

    fn compute(&self, xs: &[&NdArray], _: bool) -> NdArray
    {
        let a = xs[0];
        let b = xs[1];
        if a.shape() == &[1] {
            // a is scalar
            // unwrap is safe
            let a = NdArray::from_elem(b.shape(), *a.get(0).unwrap());
            a - b
        } else {
            a - b
        }
    }

    fn grad(&self, gy: &Tensor, _: &[&Tensor], _: &Tensor) -> Vec<Option<Tensor>>
    {
        vec![Some(gy.clone()), Some(-1 * gy)]
    }
}

impl ops::Op for MulOp {
    fn name(&self) -> &str
    {
        "MulOp"
    }

    fn compute(&self, xs: &[&NdArray], _: bool) -> NdArray
    {
        let a = xs[0];
        let b = xs[1];
        scalar_tensor_mul(a, b)
    }

    fn grad(&self, gy: &Tensor, inputs: &[&Tensor], _: &Tensor) -> Vec<Option<Tensor>>
    {
        let x0 = inputs[0];
        let x1 = inputs[1];
        vec![Some(gy * x1), Some(gy * x0)]
    }
}

impl ops::Op for DivOp {
    fn name(&self) -> &str
    {
        "DivOp"
    }

    fn compute(&self, xs: &[&NdArray], _: bool) -> NdArray
    {
        let x0 = xs[0];
        let x1 = xs[1];
        if x0.shape() == &[1] {
            // a is scalar
            // unwrap is safe
            NdArray::from_elem(x1.shape(), *x0.get(0).unwrap()) / x1
        } else {
            x0 / x1
        }
    }

    fn grad(&self, gy: &Tensor, inputs: &[&Tensor], _: &Tensor) -> Vec<Option<Tensor>>
    {
        let x0 = inputs[0];
        let x1 = inputs[1];
        let a: Tensor = gy / x1;
        let b: Tensor = -1 * x0 * ops::pow(x1, -2.) * gy;
        vec![Some(a), Some(b)]
    }
}

impl ops::Op for InplaceAddOp {
    fn inplace(&self) -> bool
    {
        true
    }

    fn name(&self) -> &str
    {
        "InplaceAddOp"
    }

    fn compute_inplace(&self, xs: &mut [&mut NdArray], _: bool)
    {
        // safe transmute
        let b: &&NdArray = unsafe { mem::transmute(&mut xs[1]) };
        let a = &mut xs[0];
        a.zip_mut_with(b, |x, &y| *x += y);
    }

    fn grad(&self, gy: &Tensor, _: &[&Tensor], _: &Tensor) -> Vec<Option<Tensor>>
    {
        vec![Some(gy.clone()), Some(gy.clone())]
    }
}

impl ops::Op for InplaceSubOp {
    fn inplace(&self) -> bool
    {
        true
    }

    fn name(&self) -> &str
    {
        "InplaceSubOp"
    }

    fn compute_inplace(&self, xs: &mut [&mut NdArray], _: bool)
    {
        // safe transmute
        let b: &&NdArray = unsafe { mem::transmute(&mut xs[1]) };
        let a = &mut xs[0];
        a.zip_mut_with(b, |x, &y| *x -= y);
    }

    fn grad(&self, gy: &Tensor, _: &[&Tensor], _: &Tensor) -> Vec<Option<Tensor>>
    {
        vec![Some(gy.clone()), Some(-1 * gy)]
    }
}

macro_rules! impl_scalar_to_tensor {
    ($func_name:ident, $scalar_type:ty) => {
        fn $func_name(val: $scalar_type) -> Tensor
        {
            Tensor(Rc::new(RawTensor {
                op: Box::new(ops::scalar::Scalar { val: val as f32 }),
                inputs: vec![],
                top_rank: 0,
            }))
        }
    };
}

// -- std::ops::{Add, Sub, Mul, Div} implementations --

macro_rules! impl_elementwise_between_tensor_and_scalar {
    (
        $trt:ident,
        $func:ident,
        $op:ident,
        $scalar_type:ty,
        $conversion_fn:ident
    ) => {

        // scalar op Tensor
        impl $trt<Tensor> for $scalar_type {
            type Output = Tensor;
            fn $func(self, rhs: Tensor) -> Self::Output
            {
                ops::apply_op($op, &[&$conversion_fn(self), &rhs])
            }
        }

        // scalar op &Tensor
        impl<'a> $trt<&'a Tensor> for $scalar_type {
            type Output = Tensor;
            fn $func(self, rhs: &'a Tensor) -> Self::Output
            {
                ops::apply_op($op, &[&$conversion_fn(self), rhs])
            }
        }

        // Tensor op scalar
        impl $trt<$scalar_type> for Tensor {
            type Output = Tensor;
            fn $func(self, rhs: $scalar_type) -> Self::Output
            {
                ops::apply_op($op, &[&self, &$conversion_fn(rhs)])
            }
        }

        // &Tensor op scalar
        impl<'a> $trt<$scalar_type> for &'a Tensor {
            type Output = Tensor;
            fn $func(self, rhs: $scalar_type) -> Self::Output
            {
                ops::apply_op($op, &[self, &$conversion_fn(rhs)])
            }
        }
    }
}

macro_rules! impl_elementwise_between_two_tensors {
    (
        $trt:ident,
        $func:ident,
        $op:ident
    ) => {
        // Tensor op Tensor
        impl $trt for Tensor {
            type Output = Tensor;
            fn $func(self, rhs: Tensor) -> Self::Output
            {
                ops::apply_op($op, &[&self, &rhs])
            }
        }

        // Tensor op &Tensor
        impl<'a> $trt<&'a Tensor> for Tensor {
            type Output = Tensor;
            fn $func(self, rhs: &Tensor) -> Self::Output
            {
                ops::apply_op($op, &[&self, rhs])
            }
        }

        // &Tensor op Tensor
        impl<'a> $trt<Tensor> for &'a Tensor {
            type Output = Tensor;
            fn $func(self, rhs: Tensor) -> Self::Output
            {
                ops::apply_op($op, &[self, &rhs])
            }
        }

        // &Tensor op &Tensor
        // lifetime of the two tensors are unrelated
        impl<'a, 'b> $trt<&'a Tensor> for &'b Tensor {
            type Output = Tensor;
            fn $func(self, rhs: &Tensor) -> Self::Output
            {
                ops::apply_op($op, &[self, rhs])
            }
        }
    };
}

impl_scalar_to_tensor!(isize_to_tensor, isize);
impl_scalar_to_tensor!(usize_to_tensor, usize);
impl_scalar_to_tensor!(i32_to_tensor, i32);
impl_scalar_to_tensor!(i64_to_tensor, i64);
impl_scalar_to_tensor!(u32_to_tensor, u32);
impl_scalar_to_tensor!(u64_to_tensor, u64);
impl_scalar_to_tensor!(f32_to_tensor, f32);
impl_scalar_to_tensor!(f64_to_tensor, f64);

impl_elementwise_between_two_tensors!(Add, add, AddOp);
impl_elementwise_between_two_tensors!(Sub, sub, SubOp);
impl_elementwise_between_two_tensors!(Mul, mul, MulOp);
impl_elementwise_between_two_tensors!(Div, div, DivOp);

impl_elementwise_between_tensor_and_scalar!(Add, add, AddOp, i32, i32_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Sub, sub, SubOp, i32, i32_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Mul, mul, MulOp, i32, i32_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Div, div, DivOp, i32, i32_to_tensor);

impl_elementwise_between_tensor_and_scalar!(Add, add, AddOp, i64, i64_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Sub, sub, SubOp, i64, i64_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Mul, mul, MulOp, i64, i64_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Div, div, DivOp, i64, i64_to_tensor);

impl_elementwise_between_tensor_and_scalar!(Add, add, AddOp, f32, f32_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Sub, sub, SubOp, f32, f32_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Mul, mul, MulOp, f32, f32_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Div, div, DivOp, f32, f32_to_tensor);

impl_elementwise_between_tensor_and_scalar!(Add, add, AddOp, f64, f64_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Sub, sub, SubOp, f64, f64_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Mul, mul, MulOp, f64, f64_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Div, div, DivOp, f64, f64_to_tensor);

impl_elementwise_between_tensor_and_scalar!(Add, add, AddOp, u32, u32_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Sub, sub, SubOp, u32, u32_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Mul, mul, MulOp, u32, u32_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Div, div, DivOp, u32, u32_to_tensor);

impl_elementwise_between_tensor_and_scalar!(Add, add, AddOp, u64, u64_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Sub, sub, SubOp, u64, u64_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Mul, mul, MulOp, u64, u64_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Div, div, DivOp, u64, u64_to_tensor);

impl_elementwise_between_tensor_and_scalar!(Add, add, AddOp, usize, usize_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Sub, sub, SubOp, usize, usize_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Mul, mul, MulOp, usize, usize_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Div, div, DivOp, usize, usize_to_tensor);

impl_elementwise_between_tensor_and_scalar!(Add, add, AddOp, isize, isize_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Sub, sub, SubOp, isize, isize_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Mul, mul, MulOp, isize, isize_to_tensor);
impl_elementwise_between_tensor_and_scalar!(Div, div, DivOp, isize, isize_to_tensor);
<|endoftext|>


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

extern crate gtk;
extern crate glib;
extern crate sysinfo;

use gtk::{Orientation, Type, Widget};
use gtk::prelude::*;

use sysinfo::*;

use std::cell::{Cell, RefCell};
use std::collections::{HashMap, HashSet};
use std::rc::Rc;

struct NoteBook {
    notebook: gtk::Notebook,
    tabs: Vec<gtk::Box>,
}

impl NoteBook {
    fn new() -> NoteBook {
        NoteBook {
            notebook: gtk::Notebook::new(),
            tabs: Vec::new(),
        }
    }

    fn create_tab(&mut self, title: &str, widget: &Widget) -> Option<u32> {
        let label = gtk::Label::new(Some(title));
        let tab = gtk::Box::new(Orientation::Horizontal, 0);

        tab.pack_start(&label, false, false, 0);
        tab.show_all();

        let index = self.notebook.append_page(widget, Some(&tab));

        self.tabs.push(tab);

        Some(index)
    }
}

#[allow(dead_code)]
struct Procs {
    left_tree: gtk::TreeView,
    scroll: gtk::ScrolledWindow,
    current_pid: Rc<Cell<Option<i64>>>,
    kill_button: gtk::Button,
    vertical_layout: gtk::Box,
    list_store: gtk::ListStore,
    columns: Vec<gtk::TreeViewColumn>,
}

impl Procs {
    pub fn new(proc_list: &HashMap<usize, Process>, note: &mut NoteBook) -> Procs {
        let left_tree = gtk::TreeView::new();
        let scroll = gtk::ScrolledWindow::new(None, None);
        let current_pid = Rc::new(Cell::new(None));
        let kill_button = gtk::Button::new_with_label("End task");
        let current_pid1 = current_pid.clone();
        let kill_button1 = kill_button.clone();

        scroll.set_min_content_height(800);
        scroll.set_min_content_width(600);

        let mut columns : Vec<gtk::TreeViewColumn> = Vec::new();

        append_column("pid", &mut columns, &left_tree);
        append_column("process name", &mut columns, &left_tree);
        append_column("cpu usage", &mut columns, &left_tree);
        append_column("memory usage (in kB)", &mut columns, &left_tree);
        // sort the name column by name_lowercase
        columns[1].set_sort_column_id(4);
        // sort the CPU column by CPU_f32
        columns[2].set_sort_column_id(5);

        let list_store = gtk::ListStore::new(&[Type::I64,       // pid
                                               Type::String,    // name
                                               Type::String,    // CPU
                                               Type::U32,       // mem
                                               Type::String,    // name_lowercase
                                               Type::F32,       // CPU_f32
                                              ]);
        for (_, pro) in proc_list {
            create_and_fill_model(&list_store, pro.pid, &pro.cmd, &pro.name, pro.cpu_usage,
                                  pro.memory);
        }

        left_tree.set_model(Some(&list_store));
        left_tree.set_headers_visible(true);
        scroll.add(&left_tree);
        let vertical_layout = gtk::Box::new(gtk::Orientation::Vertical, 0);

        left_tree.connect_cursor_changed(move |tree_view| {
            let selection = tree_view.get_selection();
            if let Some((model, iter)) = selection.get_selected() {
                let pid = Some(model.get_value(&iter, 0).get().unwrap());
                current_pid1.set(pid);
                kill_button1.set_sensitive(true);
            } else {
                current_pid1.set(None);
                kill_button1.set_sensitive(false);
            }
        });
        kill_button.set_sensitive(false);

        vertical_layout.add(&scroll);
        vertical_layout.add(&kill_button);
        let vertical_layout : Widget = vertical_layout.upcast();

        note.create_tab("Process list", &vertical_layout);
        Procs {
            left_tree: left_tree,
            scroll: scroll,
            current_pid: current_pid,
            kill_button: kill_button.clone(),
            vertical_layout: vertical_layout.downcast::<gtk::Box>().unwrap(),
            list_store: list_store,
            columns: columns,
        }
    }
}

fn append_column(title: &str, v: &mut Vec<gtk::TreeViewColumn>, left_tree: &gtk::TreeView) {
    let id = v.len() as i32;
    let renderer = gtk::CellRendererText::new();

    let column = gtk::TreeViewColumn::new();
    column.set_title(title);
    column.set_resizable(true);
    column.pack_start(&renderer, true);
    column.add_attribute(&renderer, "text", id);
    column.set_clickable(true);
    column.set_sort_column_id(id);
    left_tree.append_column(&column);
    v.push(column);
}

fn create_and_fill_model(list_store: &gtk::ListStore, pid: i64, cmdline: &str, name: &str,
                         cpu: f32, memory: u64) {
    if cmdline.len() < 1 {
        return;
    }
    list_store.insert_with_values(None,
                                  &[0, 1, 2, 3, 4, 5],
                                  &[&pid,
                                    &name,
                                    &format!("{:.1}", cpu),
                                    &memory,
                                    &name.to_lowercase(),
                                    &cpu
                                   ]);
}

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);
        }
    }
}

#[allow(dead_code)]
struct DisplaySysInfo {
    procs : Rc<RefCell<Vec<gtk::ProgressBar>>>,
    ram : gtk::ProgressBar,
    swap : gtk::ProgressBar,
    vertical_layout : gtk::Box,
}

impl DisplaySysInfo {
    pub fn new(sys1: Rc<RefCell<sysinfo::System>>, note: &mut NoteBook) -> DisplaySysInfo {
        let vertical_layout = gtk::Box::new(gtk::Orientation::Vertical, 0);
        let mut procs = Vec::new();
        let ram = gtk::ProgressBar::new();
        let swap = gtk::ProgressBar::new();

        ram.set_show_text(true);
        swap.set_show_text(true);
        vertical_layout.set_spacing(5);

        let mut i = 0;
        let mut total = false;

        vertical_layout.pack_start(&gtk::Label::new(Some("Memory usage")), false, false, 15);
        vertical_layout.add(&ram);
        vertical_layout.pack_start(&gtk::Label::new(Some("Swap usage")), false, false, 15);
        vertical_layout.add(&swap);
        vertical_layout.pack_start(&gtk::Label::new(Some("Total CPU usage")), false, false, 15);
        for pro in sys1.borrow().get_processor_list() {
            if total {
                procs.push(gtk::ProgressBar::new());
                let p : &gtk::ProgressBar = &procs[i];
                let l = gtk::Label::new(Some(&format!("{}", i)));
                let horizontal_layout = gtk::Box::new(gtk::Orientation::Horizontal, 0);

                p.set_text(Some(&format!("{:.2} %", pro.get_cpu_usage() * 100.)));
                p.set_show_text(true);
                p.set_fraction(pro.get_cpu_usage() as f64);
                horizontal_layout.pack_start(&l, false, false, 5);
                horizontal_layout.pack_start(p, true, true, 5);
                vertical_layout.add(&horizontal_layout);
            } else {
                procs.push(gtk::ProgressBar::new());
                let p : &gtk::ProgressBar = &procs[i];

                p.set_text(Some(&format!("{:.2} %", pro.get_cpu_usage() * 100.)));
                p.set_show_text(true);
                p.set_fraction(pro.get_cpu_usage() as f64);

                vertical_layout.add(p);
                vertical_layout.pack_start(&gtk::Label::new(Some("Process usage")), false,
                                           false, 15);
                total = true;
            }
            i += 1;
        }

        let vertical_layout : Widget = vertical_layout.upcast();
        note.create_tab("System usage", &vertical_layout);
        let vertical_layout : gtk::Box = vertical_layout.downcast::<gtk::Box>().unwrap();

        let mut tmp = DisplaySysInfo {
            procs: Rc::new(RefCell::new(procs)),
            ram: ram,
            swap: swap,
            vertical_layout: vertical_layout,
        };
        tmp.update_ram_display(&sys1.borrow());
        tmp
    }

    pub fn update_ram_display(&mut self, sys: &sysinfo::System) {
        let total = sys.get_total_memory();
        let used = sys.get_used_memory();
        let disp = if total < 100000 {
            format!("{} / {}KB", used, total)
        } else if total < 10000000 {
            format!("{} / {}MB", used / 1000, total / 1000)
        } else if total < 10000000000 {
            format!("{} / {}GB", used / 1000000, total / 1000000)
        } else {
            format!("{} / {}TB", used / 1000000000, total / 1000000000)
        };

        self.ram.set_text(Some(&disp));
        self.ram.set_fraction(used as f64 / total as f64);

        let total = sys.get_total_swap();
        let used = total - sys.get_used_swap();
        let disp = if total < 100000 {
            format!("{} / {}KB", used, total)
        } else if total < 10000000 {
            format!("{} / {}MB", used / 1000, total / 1000)
        } else if total < 10000000000 {
            format!("{} / {}GB", used / 1000000, total / 1000000)
        } else {
            format!("{} / {}TB", used / 1000000000, total / 1000000000)
        };

        self.swap.set_text(Some(&disp));
        self.swap.set_fraction(used as f64 / total as f64);
    }

    pub fn update_process_display(&mut self, sys: &sysinfo::System) {
        let v = &*self.procs.borrow_mut();
        let mut i = 0;

        for pro in sys.get_processor_list() {
            v[i].set_text(Some(&format!("{:.1} %", pro.get_cpu_usage() * 100.)));
            v[i].set_show_text(true);
            v[i].set_fraction(pro.get_cpu_usage() as f64);
            i += 1;
        }
    }
}

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

    let window = gtk::Window::new(gtk::WindowType::Toplevel);
    let sys = Rc::new(RefCell::new(sysinfo::System::new()));
    let mut note = NoteBook::new();
    let procs = Procs::new(sys.borrow().get_process_list(), &mut note);
    let current_pid = procs.current_pid.clone();
    let sys1 = sys.clone();

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

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

    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);

    gtk::timeout_add(1500, move || {
        // first part, deactivate sorting
        let sorted = procs.list_store.get_sort_column_id();
        procs.list_store.set_unsorted();

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

        // we re-enable the sorting
        if let Some((col, order)) = sorted {
            procs.list_store.set_sort_column_id(col, order);
        }
        glib::Continue(true)
    });

    window.add(&note.notebook);
    window.show_all();
    gtk::main();
}
<commit_msg>Don't inhibit the delete event<commit_after>#![crate_type = "bin"]

extern crate gtk;
extern crate glib;
extern crate sysinfo;

use gtk::{Orientation, Type, Widget};
use gtk::prelude::*;

use sysinfo::*;

use std::cell::{Cell, RefCell};
use std::collections::{HashMap, HashSet};
use std::rc::Rc;

struct NoteBook {
    notebook: gtk::Notebook,
    tabs: Vec<gtk::Box>,
}

impl NoteBook {
    fn new() -> NoteBook {
        NoteBook {
            notebook: gtk::Notebook::new(),
            tabs: Vec::new(),
        }
    }

    fn create_tab(&mut self, title: &str, widget: &Widget) -> Option<u32> {
        let label = gtk::Label::new(Some(title));
        let tab = gtk::Box::new(Orientation::Horizontal, 0);

        tab.pack_start(&label, false, false, 0);
        tab.show_all();

        let index = self.notebook.append_page(widget, Some(&tab));

        self.tabs.push(tab);

        Some(index)
    }
}

#[allow(dead_code)]
struct Procs {
    left_tree: gtk::TreeView,
    scroll: gtk::ScrolledWindow,
    current_pid: Rc<Cell<Option<i64>>>,
    kill_button: gtk::Button,
    vertical_layout: gtk::Box,
    list_store: gtk::ListStore,
    columns: Vec<gtk::TreeViewColumn>,
}

impl Procs {
    pub fn new(proc_list: &HashMap<usize, Process>, note: &mut NoteBook) -> Procs {
        let left_tree = gtk::TreeView::new();
        let scroll = gtk::ScrolledWindow::new(None, None);
        let current_pid = Rc::new(Cell::new(None));
        let kill_button = gtk::Button::new_with_label("End task");
        let current_pid1 = current_pid.clone();
        let kill_button1 = kill_button.clone();

        scroll.set_min_content_height(800);
        scroll.set_min_content_width(600);

        let mut columns : Vec<gtk::TreeViewColumn> = Vec::new();

        append_column("pid", &mut columns, &left_tree);
        append_column("process name", &mut columns, &left_tree);
        append_column("cpu usage", &mut columns, &left_tree);
        append_column("memory usage (in kB)", &mut columns, &left_tree);
        // sort the name column by name_lowercase
        columns[1].set_sort_column_id(4);
        // sort the CPU column by CPU_f32
        columns[2].set_sort_column_id(5);

        let list_store = gtk::ListStore::new(&[Type::I64,       // pid
                                               Type::String,    // name
                                               Type::String,    // CPU
                                               Type::U32,       // mem
                                               Type::String,    // name_lowercase
                                               Type::F32,       // CPU_f32
                                              ]);
        for (_, pro) in proc_list {
            create_and_fill_model(&list_store, pro.pid, &pro.cmd, &pro.name, pro.cpu_usage,
                                  pro.memory);
        }

        left_tree.set_model(Some(&list_store));
        left_tree.set_headers_visible(true);
        scroll.add(&left_tree);
        let vertical_layout = gtk::Box::new(gtk::Orientation::Vertical, 0);

        left_tree.connect_cursor_changed(move |tree_view| {
            let selection = tree_view.get_selection();
            if let Some((model, iter)) = selection.get_selected() {
                let pid = Some(model.get_value(&iter, 0).get().unwrap());
                current_pid1.set(pid);
                kill_button1.set_sensitive(true);
            } else {
                current_pid1.set(None);
                kill_button1.set_sensitive(false);
            }
        });
        kill_button.set_sensitive(false);

        vertical_layout.add(&scroll);
        vertical_layout.add(&kill_button);
        let vertical_layout : Widget = vertical_layout.upcast();

        note.create_tab("Process list", &vertical_layout);
        Procs {
            left_tree: left_tree,
            scroll: scroll,
            current_pid: current_pid,
            kill_button: kill_button.clone(),
            vertical_layout: vertical_layout.downcast::<gtk::Box>().unwrap(),
            list_store: list_store,
            columns: columns,
        }
    }
}

fn append_column(title: &str, v: &mut Vec<gtk::TreeViewColumn>, left_tree: &gtk::TreeView) {
    let id = v.len() as i32;
    let renderer = gtk::CellRendererText::new();

    let column = gtk::TreeViewColumn::new();
    column.set_title(title);
    column.set_resizable(true);
    column.pack_start(&renderer, true);
    column.add_attribute(&renderer, "text", id);
    column.set_clickable(true);
    column.set_sort_column_id(id);
    left_tree.append_column(&column);
    v.push(column);
}

fn create_and_fill_model(list_store: &gtk::ListStore, pid: i64, cmdline: &str, name: &str,
                         cpu: f32, memory: u64) {
    if cmdline.len() < 1 {
        return;
    }
    list_store.insert_with_values(None,
                                  &[0, 1, 2, 3, 4, 5],
                                  &[&pid,
                                    &name,
                                    &format!("{:.1}", cpu),
                                    &memory,
                                    &name.to_lowercase(),
                                    &cpu
                                   ]);
}

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);
        }
    }
}

#[allow(dead_code)]
struct DisplaySysInfo {
    procs : Rc<RefCell<Vec<gtk::ProgressBar>>>,
    ram : gtk::ProgressBar,
    swap : gtk::ProgressBar,
    vertical_layout : gtk::Box,
}

impl DisplaySysInfo {
    pub fn new(sys1: Rc<RefCell<sysinfo::System>>, note: &mut NoteBook) -> DisplaySysInfo {
        let vertical_layout = gtk::Box::new(gtk::Orientation::Vertical, 0);
        let mut procs = Vec::new();
        let ram = gtk::ProgressBar::new();
        let swap = gtk::ProgressBar::new();

        ram.set_show_text(true);
        swap.set_show_text(true);
        vertical_layout.set_spacing(5);

        let mut i = 0;
        let mut total = false;

        vertical_layout.pack_start(&gtk::Label::new(Some("Memory usage")), false, false, 15);
        vertical_layout.add(&ram);
        vertical_layout.pack_start(&gtk::Label::new(Some("Swap usage")), false, false, 15);
        vertical_layout.add(&swap);
        vertical_layout.pack_start(&gtk::Label::new(Some("Total CPU usage")), false, false, 15);
        for pro in sys1.borrow().get_processor_list() {
            if total {
                procs.push(gtk::ProgressBar::new());
                let p : &gtk::ProgressBar = &procs[i];
                let l = gtk::Label::new(Some(&format!("{}", i)));
                let horizontal_layout = gtk::Box::new(gtk::Orientation::Horizontal, 0);

                p.set_text(Some(&format!("{:.2} %", pro.get_cpu_usage() * 100.)));
                p.set_show_text(true);
                p.set_fraction(pro.get_cpu_usage() as f64);
                horizontal_layout.pack_start(&l, false, false, 5);
                horizontal_layout.pack_start(p, true, true, 5);
                vertical_layout.add(&horizontal_layout);
            } else {
                procs.push(gtk::ProgressBar::new());
                let p : &gtk::ProgressBar = &procs[i];

                p.set_text(Some(&format!("{:.2} %", pro.get_cpu_usage() * 100.)));
                p.set_show_text(true);
                p.set_fraction(pro.get_cpu_usage() as f64);

                vertical_layout.add(p);
                vertical_layout.pack_start(&gtk::Label::new(Some("Process usage")), false,
                                           false, 15);
                total = true;
            }
            i += 1;
        }

        let vertical_layout : Widget = vertical_layout.upcast();
        note.create_tab("System usage", &vertical_layout);
        let vertical_layout : gtk::Box = vertical_layout.downcast::<gtk::Box>().unwrap();

        let mut tmp = DisplaySysInfo {
            procs: Rc::new(RefCell::new(procs)),
            ram: ram,
            swap: swap,
            vertical_layout: vertical_layout,
        };
        tmp.update_ram_display(&sys1.borrow());
        tmp
    }

    pub fn update_ram_display(&mut self, sys: &sysinfo::System) {
        let total = sys.get_total_memory();
        let used = sys.get_used_memory();
        let disp = if total < 100000 {
            format!("{} / {}KB", used, total)
        } else if total < 10000000 {
            format!("{} / {}MB", used / 1000, total / 1000)
        } else if total < 10000000000 {
            format!("{} / {}GB", used / 1000000, total / 1000000)
        } else {
            format!("{} / {}TB", used / 1000000000, total / 1000000000)
        };

        self.ram.set_text(Some(&disp));
        self.ram.set_fraction(used as f64 / total as f64);

        let total = sys.get_total_swap();
        let used = total - sys.get_used_swap();
        let disp = if total < 100000 {
            format!("{} / {}KB", used, total)
        } else if total < 10000000 {
            format!("{} / {}MB", used / 1000, total / 1000)
        } else if total < 10000000000 {
            format!("{} / {}GB", used / 1000000, total / 1000000)
        } else {
            format!("{} / {}TB", used / 1000000000, total / 1000000000)
        };

        self.swap.set_text(Some(&disp));
        self.swap.set_fraction(used as f64 / total as f64);
    }

    pub fn update_process_display(&mut self, sys: &sysinfo::System) {
        let v = &*self.procs.borrow_mut();
        let mut i = 0;

        for pro in sys.get_processor_list() {
            v[i].set_text(Some(&format!("{:.1} %", pro.get_cpu_usage() * 100.)));
            v[i].set_show_text(true);
            v[i].set_fraction(pro.get_cpu_usage() as f64);
            i += 1;
        }
    }
}

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

    let window = gtk::Window::new(gtk::WindowType::Toplevel);
    let sys = Rc::new(RefCell::new(sysinfo::System::new()));
    let mut note = NoteBook::new();
    let procs = Procs::new(sys.borrow().get_process_list(), &mut note);
    let current_pid = procs.current_pid.clone();
    let sys1 = sys.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);

    gtk::timeout_add(1500, move || {
        // first part, deactivate sorting
        let sorted = procs.list_store.get_sort_column_id();
        procs.list_store.set_unsorted();

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

        // we re-enable the sorting
        if let Some((col, order)) = sorted {
            procs.list_store.set_sort_column_id(col, order);
        }
        glib::Continue(true)
    });

    window.add(&note.notebook);
    window.show_all();
    gtk::main();
}
<|endoftext|>


<commit_before>// TODO rm allow(unused)
#![allow(unused)]
#![allow(unsafe_code)]

use std::{
    convert::{TryFrom, TryInto},
    fmt::{self, Debug},
    fs::File,
    mem::{transmute, MaybeUninit},
    path::Path,
    sync::{
        atomic::{AtomicU32, Ordering::Acquire},
        Arc,
    },
};

use crossbeam_epoch::pin;

use crate::{
    pagecache::{pread_exact, pwrite_all, MessageKind},
    stack::Stack,
    Error, Lsn, Result,
};

#[cfg(not(feature = "testing"))]
const MIN_SZ: u64 = 64 * 1024;

#[cfg(feature = "testing")]
const MIN_SZ: u64 = 32;

const MIN_TRAILING_ZEROS: u64 = MIN_SZ.trailing_zeros() as u64;

pub type SlabId = u8;
pub type SlabIdx = u32;

/// A unique identifier for a particular slot in the heap
#[allow(clippy::module_name_repetitions)]
#[derive(Clone, Copy, PartialOrd, Ord, Eq, PartialEq, Hash)]
pub struct HeapId {
    pub location: u64,
    pub original_lsn: Lsn,
}

impl Debug for HeapId {
    fn fmt(
        &self,
        f: &mut fmt::Formatter<'_>,
    ) -> std::result::Result<(), fmt::Error> {
        let (slab, idx, original_lsn) = self.decompose();
        f.debug_struct("HeapId")
            .field("slab", &slab)
            .field("idx", &idx)
            .field("original_lsn", &original_lsn)
            .finish()
    }
}

impl HeapId {
    pub fn decompose(&self) -> (SlabId, SlabIdx, Lsn) {
        const IDX_MASK: u64 = (1 << 32) - 1;
        let slab_id =
            u8::try_from((self.location >> 32).trailing_zeros()).unwrap();
        let slab_idx = u32::try_from(self.location & IDX_MASK).unwrap();
        (slab_id, slab_idx, self.original_lsn)
    }

    pub fn compose(
        slab_id: SlabId,
        slab_idx: SlabIdx,
        original_lsn: Lsn,
    ) -> HeapId {
        let slab = 1 << (32 + u64::from(slab_id));
        let heap_id = slab | u64::from(slab_idx);
        HeapId { location: heap_id, original_lsn }
    }

    fn offset(&self) -> u64 {
        let (slab_id, idx, _) = self.decompose();
        slab_id_to_size(slab_id) * u64::from(idx)
    }

    fn slab_size(&self) -> u64 {
        let (slab_id, idx, _) = self.decompose();
        slab_id_to_size(slab_id)
    }
}

pub(crate) fn slab_size(size: u64) -> u64 {
    slab_id_to_size(size_to_slab_id(size))
}

fn slab_id_to_size(slab_id: u8) -> u64 {
    1 << (MIN_TRAILING_ZEROS + u64::from(slab_id))
}

fn size_to_slab_id(size: u64) -> SlabId {
    // find the power of 2 that is at least 64k
    let normalized_size = std::cmp::max(MIN_SZ, size.next_power_of_two());

    // drop the lowest unused bits
    let rebased_size = normalized_size >> MIN_TRAILING_ZEROS;

    u8::try_from(rebased_size.trailing_zeros()).unwrap()
}

pub(crate) struct Reservation {
    slab_free: Arc<Stack<u32>>,
    completed: bool,
    file: File,
    pub heap_id: HeapId,
}

impl Drop for Reservation {
    fn drop(&mut self) {
        if !self.completed {
            let (slab_id, idx, _) = self.heap_id.decompose();
            self.slab_free.push(idx, &pin());
        }
    }
}

impl Reservation {
    pub fn complete(mut self, data: &[u8]) -> Result<HeapId> {
        log::trace!(
            "Heap::complete({:?}) to offset {} in file {:?}",
            self.heap_id,
            self.heap_id.offset(),
            self.file
        );
        assert_eq!(data.len() as u64, self.heap_id.slab_size());

        pwrite_all(&self.file, data, self.heap_id.offset())?;
        self.file.sync_all()?;

        // if this is not reached due to an IO error,
        // the offset will be returned to the Slab in Drop
        self.completed = true;

        Ok(self.heap_id)
    }

    pub fn abort(self) {
        // actual logic in Drop
    }
}

#[derive(Debug)]
pub(crate) struct Heap {
    // each slab stores
    // items that are double
    // the size of the previous,
    // ranging from 64k in the
    // smallest slab to 2^48 in
    // the last.
    slabs: [Slab; 32],
}

impl Heap {
    pub fn start<P: AsRef<Path>>(p: P) -> Result<Heap> {
        let mut slabs: [MaybeUninit<Slab>; 32] = unsafe { std::mem::zeroed() };

        for slab_id in 0..32 {
            let slab = Slab::start(&p, slab_id)?;
            slabs[slab_id as usize] = MaybeUninit::new(slab);
        }

        Ok(Heap { slabs: unsafe { transmute(slabs) } })
    }

    pub fn gc_unknown_items(
        &self,
        _snapshot: &crate::pagecache::Snapshot,
    ) -> Result<()> {
        //TODO todo!()
        Ok(())
    }

    pub fn read(
        &self,
        heap_id: HeapId,
        use_compression: bool,
    ) -> Result<(MessageKind, Vec<u8>)> {
        log::trace!("Heap::read({:?})", heap_id);
        let (slab_id, slab_idx, original_lsn) = heap_id.decompose();
        self.slabs[slab_id as usize].read(
            slab_idx,
            original_lsn,
            use_compression,
        )
    }

    pub fn free(&self, heap_id: HeapId) {
        log::trace!("Heap::free({:?})", heap_id);
        let (slab_id, slab_idx, _) = heap_id.decompose();
        self.slabs[slab_id as usize].free(slab_idx)
    }

    pub fn reserve(&self, size: u64, original_lsn: Lsn) -> Reservation {
        assert!(size < 1 << 48);
        let slab_id = size_to_slab_id(size);
        let ret = self.slabs[slab_id as usize].reserve(size, original_lsn);
        log::trace!("Heap::reserve({}) -> {:?}", size, ret.heap_id);
        ret
    }
}

#[derive(Debug)]
struct Slab {
    file: File,
    slab_id: u8,
    tip: AtomicU32,
    free: Arc<Stack<u32>>,
}

impl Slab {
    pub fn start<P: AsRef<Path>>(directory: P, slab_id: u8) -> Result<Slab> {
        let bs = slab_id_to_size(slab_id);
        let free = Arc::new(Stack::default());

        let mut options = std::fs::OpenOptions::new();
        options.create(true);
        options.read(true);
        options.write(true);

        let file =
            options.open(directory.as_ref().join(format!("{}", slab_id)))?;
        let len = file.metadata()?.len();
        let max_idx = len / bs;
        log::trace!(
            "starting heap slab for sizes of {}. tip: {} max idx: {}",
            bs,
            len,
            max_idx
        );
        let tip = AtomicU32::new(u32::try_from(max_idx).unwrap());

        Ok(Slab { file, slab_id, tip, free })
    }

    fn read(
        &self,
        slab_idx: SlabIdx,
        original_lsn: Lsn,
        use_compression: bool,
    ) -> Result<(MessageKind, Vec<u8>)> {
        let bs = slab_id_to_size(self.slab_id);
        let offset = u64::from(slab_idx) * bs;

        log::trace!("reading heap slab slot {} at offset {}", slab_idx, offset);

        let mut heap_buf = vec![0; usize::try_from(bs).unwrap()];

        pread_exact(&self.file, &mut heap_buf, offset)?;

        let stored_crc =
            u32::from_le_bytes(heap_buf[1..5].as_ref().try_into().unwrap());

        let mut hasher = crc32fast::Hasher::new();
        hasher.update(&heap_buf[0..1]);
        hasher.update(&heap_buf[5..]);
        let actual_crc = hasher.finalize();

        if actual_crc == stored_crc {
            let actual_lsn = Lsn::from_le_bytes(
                heap_buf[5..13].as_ref().try_into().unwrap(),
            );
            if actual_lsn != original_lsn {
                log::error!(
                    "heap slot lsn {} does not match expected original lsn {}",
                    actual_lsn,
                    original_lsn
                );
                return Err(Error::corruption(None));
            }
            let buf = heap_buf[13..].to_vec();
            let buf = if use_compression {
                crate::pagecache::decompress(buf)
            } else {
                buf
            };
            Ok((MessageKind::from(heap_buf[0]), buf))
        } else {
            log::error!(
                "heap message CRC does not match contents. stored: {} actual: {}",
                stored_crc,
                actual_crc
            );
            Err(Error::corruption(None))
        }
    }

    fn reserve(&self, size: u64, original_lsn: Lsn) -> Reservation {
        let idx = if let Some(idx) = self.free.pop(&pin()) {
            log::trace!(
                "reusing heap index {} in slab for sizes of {}",
                idx,
                slab_id_to_size(self.slab_id),
            );
            idx
        } else {
            log::trace!(
                "no free heap slots in slab for sizes of {}",
                slab_id_to_size(self.slab_id),
            );
            self.tip.fetch_add(1, Acquire)
        };

        log::trace!(
            "heap reservation for slot {} in the slab for sizes of {}",
            idx,
            slab_id_to_size(self.slab_id),
        );

        let heap_id = HeapId::compose(self.slab_id, idx, original_lsn);

        Reservation {
            slab_free: self.free.clone(),
            completed: false,
            file: self.file.try_clone().unwrap(),
            heap_id,
        }
    }

    fn free(&self, idx: u32) {
        self.punch_hole(idx);
        self.free.push(idx, &pin());
    }

    fn punch_hole(&self, idx: u32) {
        let bs = i64::try_from(slab_id_to_size(self.slab_id)).unwrap();
        let offset = i64::from(idx) * bs;

        #[cfg(target_os = "linux")]
        {
            use std::{
                os::unix::io::AsRawFd,
                sync::atomic::{AtomicBool, Ordering::Relaxed},
            };

            use libc::{fallocate, FALLOC_FL_KEEP_SIZE, FALLOC_FL_PUNCH_HOLE};

            static HOLE_PUNCHING_ENABLED: AtomicBool = AtomicBool::new(false);

            if HOLE_PUNCHING_ENABLED.load(Relaxed) {
                let mode = FALLOC_FL_KEEP_SIZE | FALLOC_FL_PUNCH_HOLE;

                let fd = self.file.as_raw_fd();

                let ret = unsafe { fallocate(fd, mode, offset, bs) };

                if ret != 0 {
                    let err = std::io::Error::last_os_error();
                    log::error!(
                        "failed to punch hole in heap file: {:?}. disabling hole punching",
                        err
                    );
                    HOLE_PUNCHING_ENABLED.store(false, Relaxed);
                }
            }
        }
    }
}
<commit_msg>Properly handle 32-bit systems for hole punching<commit_after>// TODO rm allow(unused)
#![allow(unused)]
#![allow(unsafe_code)]

use std::{
    convert::{TryFrom, TryInto},
    fmt::{self, Debug},
    fs::File,
    mem::{transmute, MaybeUninit},
    path::Path,
    sync::{
        atomic::{AtomicU32, Ordering::Acquire},
        Arc,
    },
};

use crossbeam_epoch::pin;

use crate::{
    pagecache::{pread_exact, pwrite_all, MessageKind},
    stack::Stack,
    Error, Lsn, Result,
};

#[cfg(not(feature = "testing"))]
const MIN_SZ: u64 = 64 * 1024;

#[cfg(feature = "testing")]
const MIN_SZ: u64 = 32;

const MIN_TRAILING_ZEROS: u64 = MIN_SZ.trailing_zeros() as u64;

pub type SlabId = u8;
pub type SlabIdx = u32;

/// A unique identifier for a particular slot in the heap
#[allow(clippy::module_name_repetitions)]
#[derive(Clone, Copy, PartialOrd, Ord, Eq, PartialEq, Hash)]
pub struct HeapId {
    pub location: u64,
    pub original_lsn: Lsn,
}

impl Debug for HeapId {
    fn fmt(
        &self,
        f: &mut fmt::Formatter<'_>,
    ) -> std::result::Result<(), fmt::Error> {
        let (slab, idx, original_lsn) = self.decompose();
        f.debug_struct("HeapId")
            .field("slab", &slab)
            .field("idx", &idx)
            .field("original_lsn", &original_lsn)
            .finish()
    }
}

impl HeapId {
    pub fn decompose(&self) -> (SlabId, SlabIdx, Lsn) {
        const IDX_MASK: u64 = (1 << 32) - 1;
        let slab_id =
            u8::try_from((self.location >> 32).trailing_zeros()).unwrap();
        let slab_idx = u32::try_from(self.location & IDX_MASK).unwrap();
        (slab_id, slab_idx, self.original_lsn)
    }

    pub fn compose(
        slab_id: SlabId,
        slab_idx: SlabIdx,
        original_lsn: Lsn,
    ) -> HeapId {
        let slab = 1 << (32 + u64::from(slab_id));
        let heap_id = slab | u64::from(slab_idx);
        HeapId { location: heap_id, original_lsn }
    }

    fn offset(&self) -> u64 {
        let (slab_id, idx, _) = self.decompose();
        slab_id_to_size(slab_id) * u64::from(idx)
    }

    fn slab_size(&self) -> u64 {
        let (slab_id, idx, _) = self.decompose();
        slab_id_to_size(slab_id)
    }
}

pub(crate) fn slab_size(size: u64) -> u64 {
    slab_id_to_size(size_to_slab_id(size))
}

fn slab_id_to_size(slab_id: u8) -> u64 {
    1 << (MIN_TRAILING_ZEROS + u64::from(slab_id))
}

fn size_to_slab_id(size: u64) -> SlabId {
    // find the power of 2 that is at least 64k
    let normalized_size = std::cmp::max(MIN_SZ, size.next_power_of_two());

    // drop the lowest unused bits
    let rebased_size = normalized_size >> MIN_TRAILING_ZEROS;

    u8::try_from(rebased_size.trailing_zeros()).unwrap()
}

pub(crate) struct Reservation {
    slab_free: Arc<Stack<u32>>,
    completed: bool,
    file: File,
    pub heap_id: HeapId,
}

impl Drop for Reservation {
    fn drop(&mut self) {
        if !self.completed {
            let (slab_id, idx, _) = self.heap_id.decompose();
            self.slab_free.push(idx, &pin());
        }
    }
}

impl Reservation {
    pub fn complete(mut self, data: &[u8]) -> Result<HeapId> {
        log::trace!(
            "Heap::complete({:?}) to offset {} in file {:?}",
            self.heap_id,
            self.heap_id.offset(),
            self.file
        );
        assert_eq!(data.len() as u64, self.heap_id.slab_size());

        pwrite_all(&self.file, data, self.heap_id.offset())?;
        self.file.sync_all()?;

        // if this is not reached due to an IO error,
        // the offset will be returned to the Slab in Drop
        self.completed = true;

        Ok(self.heap_id)
    }

    pub fn abort(self) {
        // actual logic in Drop
    }
}

#[derive(Debug)]
pub(crate) struct Heap {
    // each slab stores
    // items that are double
    // the size of the previous,
    // ranging from 64k in the
    // smallest slab to 2^48 in
    // the last.
    slabs: [Slab; 32],
}

impl Heap {
    pub fn start<P: AsRef<Path>>(p: P) -> Result<Heap> {
        let mut slabs: [MaybeUninit<Slab>; 32] = unsafe { std::mem::zeroed() };

        for slab_id in 0..32 {
            let slab = Slab::start(&p, slab_id)?;
            slabs[slab_id as usize] = MaybeUninit::new(slab);
        }

        Ok(Heap { slabs: unsafe { transmute(slabs) } })
    }

    pub fn gc_unknown_items(
        &self,
        _snapshot: &crate::pagecache::Snapshot,
    ) -> Result<()> {
        //TODO todo!()
        Ok(())
    }

    pub fn read(
        &self,
        heap_id: HeapId,
        use_compression: bool,
    ) -> Result<(MessageKind, Vec<u8>)> {
        log::trace!("Heap::read({:?})", heap_id);
        let (slab_id, slab_idx, original_lsn) = heap_id.decompose();
        self.slabs[slab_id as usize].read(
            slab_idx,
            original_lsn,
            use_compression,
        )
    }

    pub fn free(&self, heap_id: HeapId) {
        log::trace!("Heap::free({:?})", heap_id);
        let (slab_id, slab_idx, _) = heap_id.decompose();
        self.slabs[slab_id as usize].free(slab_idx)
    }

    pub fn reserve(&self, size: u64, original_lsn: Lsn) -> Reservation {
        assert!(size < 1 << 48);
        let slab_id = size_to_slab_id(size);
        let ret = self.slabs[slab_id as usize].reserve(size, original_lsn);
        log::trace!("Heap::reserve({}) -> {:?}", size, ret.heap_id);
        ret
    }
}

#[derive(Debug)]
struct Slab {
    file: File,
    slab_id: u8,
    tip: AtomicU32,
    free: Arc<Stack<u32>>,
}

impl Slab {
    pub fn start<P: AsRef<Path>>(directory: P, slab_id: u8) -> Result<Slab> {
        let bs = slab_id_to_size(slab_id);
        let free = Arc::new(Stack::default());

        let mut options = std::fs::OpenOptions::new();
        options.create(true);
        options.read(true);
        options.write(true);

        let file =
            options.open(directory.as_ref().join(format!("{}", slab_id)))?;
        let len = file.metadata()?.len();
        let max_idx = len / bs;
        log::trace!(
            "starting heap slab for sizes of {}. tip: {} max idx: {}",
            bs,
            len,
            max_idx
        );
        let tip = AtomicU32::new(u32::try_from(max_idx).unwrap());

        Ok(Slab { file, slab_id, tip, free })
    }

    fn read(
        &self,
        slab_idx: SlabIdx,
        original_lsn: Lsn,
        use_compression: bool,
    ) -> Result<(MessageKind, Vec<u8>)> {
        let bs = slab_id_to_size(self.slab_id);
        let offset = u64::from(slab_idx) * bs;

        log::trace!("reading heap slab slot {} at offset {}", slab_idx, offset);

        let mut heap_buf = vec![0; usize::try_from(bs).unwrap()];

        pread_exact(&self.file, &mut heap_buf, offset)?;

        let stored_crc =
            u32::from_le_bytes(heap_buf[1..5].as_ref().try_into().unwrap());

        let mut hasher = crc32fast::Hasher::new();
        hasher.update(&heap_buf[0..1]);
        hasher.update(&heap_buf[5..]);
        let actual_crc = hasher.finalize();

        if actual_crc == stored_crc {
            let actual_lsn = Lsn::from_le_bytes(
                heap_buf[5..13].as_ref().try_into().unwrap(),
            );
            if actual_lsn != original_lsn {
                log::error!(
                    "heap slot lsn {} does not match expected original lsn {}",
                    actual_lsn,
                    original_lsn
                );
                return Err(Error::corruption(None));
            }
            let buf = heap_buf[13..].to_vec();
            let buf = if use_compression {
                crate::pagecache::decompress(buf)
            } else {
                buf
            };
            Ok((MessageKind::from(heap_buf[0]), buf))
        } else {
            log::error!(
                "heap message CRC does not match contents. stored: {} actual: {}",
                stored_crc,
                actual_crc
            );
            Err(Error::corruption(None))
        }
    }

    fn reserve(&self, size: u64, original_lsn: Lsn) -> Reservation {
        let idx = if let Some(idx) = self.free.pop(&pin()) {
            log::trace!(
                "reusing heap index {} in slab for sizes of {}",
                idx,
                slab_id_to_size(self.slab_id),
            );
            idx
        } else {
            log::trace!(
                "no free heap slots in slab for sizes of {}",
                slab_id_to_size(self.slab_id),
            );
            self.tip.fetch_add(1, Acquire)
        };

        log::trace!(
            "heap reservation for slot {} in the slab for sizes of {}",
            idx,
            slab_id_to_size(self.slab_id),
        );

        let heap_id = HeapId::compose(self.slab_id, idx, original_lsn);

        Reservation {
            slab_free: self.free.clone(),
            completed: false,
            file: self.file.try_clone().unwrap(),
            heap_id,
        }
    }

    fn free(&self, idx: u32) {
        self.punch_hole(idx);
        self.free.push(idx, &pin());
    }

    fn punch_hole(&self, idx: u32) {
        #[cfg(target_os = "linux")]
        {
            use std::{
                os::unix::io::AsRawFd,
                sync::atomic::{AtomicBool, Ordering::Relaxed},
            };

            use libc::{fallocate, FALLOC_FL_KEEP_SIZE, FALLOC_FL_PUNCH_HOLE};

            static HOLE_PUNCHING_ENABLED: AtomicBool = AtomicBool::new(false);

            #[cfg(target_pointer_width = "64")]
            let bs = i64::try_from(slab_id_to_size(self.slab_id)).unwrap();

            #[cfg(target_pointer_width = "32")]
            let bs = i32::try_from(slab_id_to_size(self.slab_id)).unwrap();

            #[cfg(target_pointer_width = "64")]
            let offset = i64::from(idx) * bs;

            #[cfg(target_pointer_width = "32")]
            let offset = i32::try_from(idx).unwrap() * bs;

            if HOLE_PUNCHING_ENABLED.load(Relaxed) {
                let mode = FALLOC_FL_KEEP_SIZE | FALLOC_FL_PUNCH_HOLE;

                let fd = self.file.as_raw_fd();

                let ret = unsafe { fallocate(fd, mode, offset, bs) };

                if ret != 0 {
                    let err = std::io::Error::last_os_error();
                    log::error!(
                        "failed to punch hole in heap file: {:?}. disabling hole punching",
                        err
                    );
                    HOLE_PUNCHING_ENABLED.store(false, Relaxed);
                }
            }
        }
    }
}
<|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};

#[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!();

    // 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 pre = match doc["pre"].as_str() {
            Some(x) => Some(x.to_string()),
            None    => None
        };

        // A clojure used to capture the current local root and pre Options.
        // This way we can call the clojure to create common SendKeys command
        // like changing the directory or executing a system command from the
        // `pre` option.
        let common_commands = |target: String| -> Vec<Command> {
            let mut commands2 = vec!();

            // SendKeys to change to the `root` directory
            if root.is_some() {
                let r = root.clone().unwrap();
                commands2.push(Command::SendKeys(SendKeys{
                    target: target.clone(),
                    exec: format!("cd {}", r)
                }));
            };

            // SendKeys for the Pre option
            if pre.is_some() {
                let p = pre.clone().unwrap();
                commands2.push(Command::SendKeys(SendKeys{
                    target: target.clone(),
                    exec: p
                }));
            };

            commands2
        };

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

        for window in windows.iter() {
            match window {
                &Yaml::Hash(ref h)  => {
                    for (k, v) in h {
                        if v.as_hash().is_some() {
                            commands.push(Command::Window(Window{
                                session_name: project_name.clone(),
                                name: k.as_str().unwrap().to_string()
                            }));

                            let target = format!("{}:{}", project_name, k.as_str().unwrap());
                            commands.append(&mut try!(pane_matcher(v, &target, &common_commands)));
                        } else {
                            commands.push(Command::Window(Window{
                                session_name: project_name.clone(),
                                name: try!(k.as_str().ok_or_else(|| "Windows require being named in your config.").map(|x| x.to_string()))
                            }));

                            let t = format!("{}:{}", project_name, k.as_str().unwrap()).to_string();
                            commands.append(&mut common_commands(t.to_string()));

                            // SendKeys for the exec command
                            if v.as_str().is_some() {
                                let ex = v.as_str().unwrap();
                                if !ex.is_empty() {
                                    commands.push(Command::SendKeys(SendKeys{
                                        target: format!("{}:{}", project_name, k.as_str().unwrap()).to_string(),
                                        exec: v.as_str().unwrap().to_string()
                                    }));
                                };
                            }
                        }
                    }
                },
                &Yaml::String(ref s) => {
                    commands.push(Command::Window(Window{
                        session_name: project_name.clone(),
                        name: s.clone()
                    }));

                    let t = format!("{}:{}", &project_name, &s);
                    commands.append(&mut common_commands(t.to_string()));
                },
                &Yaml::Integer(ref s) => {
                    commands.push(Command::Window(Window{
                        session_name: project_name.clone(),
                        name: s.to_string()
                    }));

                    let t = format!("{}:{}", &project_name, &s);
                    commands.append(&mut common_commands(t.to_string()));
                },
                _ => panic!("Muxed config file formatting isn't recognized.")
            };
        };
    };

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

    let (first, commands) = commands.split_first().unwrap();
    let mut remains = commands.to_vec();

    match first {
        &Command::Window(ref w) => {
            remains.insert(0, Command::Session(Session{
                name: project_name.clone(),
                window_name: w.name.clone()
            }));
        },
        _ => {}
    };

    Ok(remains)
}

/// Pane matcher is for breaking apart the panes. Splitting windows when needed
/// and executing commands as needed.
fn pane_matcher<T>(window: &Yaml, target: &str, common_commands: T) -> Result<Vec<Command>, String>
    where T : Fn(String) -> Vec<Command> {

    let mut commands = vec!();
    let panes = window["panes"].as_vec().expect("Something is wrong with panes.");

    for (i, pane) in panes.iter().enumerate() {
        let t = format!("{}.{}", target, i);
        // 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: t.to_string()
            }));
        };

        // Call the common_commands clojure to execute `cd` and `pre` options in
        // pane splits.
        commands.append(&mut common_commands(t.to_string()));

        // Execute given commands in each new pane after all splits are
        // complete.
        if pane.as_str().is_some() {
            let p = pane.as_str().unwrap();
            if !p.is_empty() {
                commands.push(Command::SendKeys(SendKeys{
                    target: t.to_string(),
                    exec: p.to_string()
                }));
            };
        };
    };

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

    Ok(commands)
}

#[test]
pub fn expect_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 expect_2_windows_from_array() {
    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(), 2)
}

#[test]
pub fn expect_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 expect_2_windows_with_mixed_type_names() {
    let s = "---
windows: [1, 'vim', 3]
";
    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(), 2)
}

#[test]
pub fn expect_2_windows_from_list() {
    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(), 2)
}

#[test]
pub fn expect_ok_with_empty_syscommands() {
    let s = "---
windows:
  - editor:
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let result = main(&yaml, &"muxed".to_string(), false);
    assert!(result.is_ok())
}

#[test]
pub fn expect_no_send_keys_commands() {
    let s = "---
windows:
  - editor:
";

    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::SendKeys(_) => true,
        _ => false
    }).collect();

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

#[test]
pub fn expect_err_with_nameless_window() {
    let s = "---
windows:
  - : ls
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let result = main(&yaml, &"muxed".to_string(), false);
    assert!(result.is_err())
}

#[test]
pub fn expect_ok_with_empty_panes_syscommands() {
    let s = "---
windows:
  - cargo:
      layout: 'main-vertical'
      panes:
        -
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let result = main(&yaml, &"muxed".to_string(), false);
    assert!(result.is_ok())
}

#[test]
pub fn expect_no_send_keys_with_empty_panes_syscommands() {
    let s = "---
windows:
  - editor:
      panes:
        -
";

    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::SendKeys(_) => true,
        _ => false
    }).collect();

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

#[test]
pub fn expect_1_split_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::Split(_) => true,
        _ => false
    }).collect();

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

#[test]
pub fn expect_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 expect_1_session_with_panes_array() {
    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)
}

#[test]
pub fn expect_no_layout() {
    let s = "---
windows:
  - editor:
      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(), 0)
}

#[test]
pub fn expect_three_send_keys_commands_from_pre() {
    let s = "---
pre: 'ls'
windows:
  - editor:
      panes:
        -
        -
  - logs:
";
    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::SendKeys(_) => true,
        _ => false
    }).collect();

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

#[test]
pub fn expect_no_send_keys_with_blank_panes() {
    let s = "---
windows:
  - editor:
      panes: ['','','']
";

    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::SendKeys(_) => true,
        _ => false
    }).collect();

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

#[test]
pub fn expect_no_send_keys_with_blank_window() {
    let s = "---
windows:
  - editor: ''
";

    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::SendKeys(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 0)
}
<commit_msg>Issue #11 Quote the root directory<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};

#[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!();

    // 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 pre = match doc["pre"].as_str() {
            Some(x) => Some(x.to_string()),
            None    => None
        };

        // A clojure used to capture the current local root and pre Options.
        // This way we can call the clojure to create common SendKeys command
        // like changing the directory or executing a system command from the
        // `pre` option.
        let common_commands = |target: String| -> Vec<Command> {
            let mut commands2 = vec!();

            // SendKeys to change to the `root` directory
            if root.is_some() {
                let r = root.clone().unwrap();
                commands2.push(Command::SendKeys(SendKeys{
                    target: target.clone(),
                    exec: format!("cd \"{}\"", r)
                }));
            };

            // SendKeys for the Pre option
            if pre.is_some() {
                let p = pre.clone().unwrap();
                commands2.push(Command::SendKeys(SendKeys{
                    target: target.clone(),
                    exec: p
                }));
            };

            commands2
        };

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

        for window in windows.iter() {
            match window {
                &Yaml::Hash(ref h)  => {
                    for (k, v) in h {
                        if v.as_hash().is_some() {
                            commands.push(Command::Window(Window{
                                session_name: project_name.clone(),
                                name: k.as_str().unwrap().to_string()
                            }));

                            let target = format!("{}:{}", project_name, k.as_str().unwrap());
                            commands.append(&mut try!(pane_matcher(v, &target, &common_commands)));
                        } else {
                            commands.push(Command::Window(Window{
                                session_name: project_name.clone(),
                                name: try!(k.as_str().ok_or_else(|| "Windows require being named in your config.").map(|x| x.to_string()))
                            }));

                            let t = format!("{}:{}", project_name, k.as_str().unwrap()).to_string();
                            commands.append(&mut common_commands(t.to_string()));

                            // SendKeys for the exec command
                            if v.as_str().is_some() {
                                let ex = v.as_str().unwrap();
                                if !ex.is_empty() {
                                    commands.push(Command::SendKeys(SendKeys{
                                        target: format!("{}:{}", project_name, k.as_str().unwrap()).to_string(),
                                        exec: v.as_str().unwrap().to_string()
                                    }));
                                };
                            }
                        }
                    }
                },
                &Yaml::String(ref s) => {
                    commands.push(Command::Window(Window{
                        session_name: project_name.clone(),
                        name: s.clone()
                    }));

                    let t = format!("{}:{}", &project_name, &s);
                    commands.append(&mut common_commands(t.to_string()));
                },
                &Yaml::Integer(ref s) => {
                    commands.push(Command::Window(Window{
                        session_name: project_name.clone(),
                        name: s.to_string()
                    }));

                    let t = format!("{}:{}", &project_name, &s);
                    commands.append(&mut common_commands(t.to_string()));
                },
                _ => panic!("Muxed config file formatting isn't recognized.")
            };
        };
    };

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

    let (first, commands) = commands.split_first().unwrap();
    let mut remains = commands.to_vec();

    match first {
        &Command::Window(ref w) => {
            remains.insert(0, Command::Session(Session{
                name: project_name.clone(),
                window_name: w.name.clone()
            }));
        },
        _ => {}
    };

    Ok(remains)
}

/// Pane matcher is for breaking apart the panes. Splitting windows when needed
/// and executing commands as needed.
fn pane_matcher<T>(window: &Yaml, target: &str, common_commands: T) -> Result<Vec<Command>, String>
    where T : Fn(String) -> Vec<Command> {

    let mut commands = vec!();
    let panes = window["panes"].as_vec().expect("Something is wrong with panes.");

    for (i, pane) in panes.iter().enumerate() {
        let t = format!("{}.{}", target, i);
        // 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: t.to_string()
            }));
        };

        // Call the common_commands clojure to execute `cd` and `pre` options in
        // pane splits.
        commands.append(&mut common_commands(t.to_string()));

        // Execute given commands in each new pane after all splits are
        // complete.
        if pane.as_str().is_some() {
            let p = pane.as_str().unwrap();
            if !p.is_empty() {
                commands.push(Command::SendKeys(SendKeys{
                    target: t.to_string(),
                    exec: p.to_string()
                }));
            };
        };
    };

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

    Ok(commands)
}

#[test]
pub fn expect_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 expect_2_windows_from_array() {
    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(), 2)
}

#[test]
pub fn expect_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 expect_2_windows_with_mixed_type_names() {
    let s = "---
windows: [1, 'vim', 3]
";
    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(), 2)
}

#[test]
pub fn expect_2_windows_from_list() {
    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(), 2)
}

#[test]
pub fn expect_ok_with_empty_syscommands() {
    let s = "---
windows:
  - editor:
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let result = main(&yaml, &"muxed".to_string(), false);
    assert!(result.is_ok())
}

#[test]
pub fn expect_no_send_keys_commands() {
    let s = "---
windows:
  - editor:
";

    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::SendKeys(_) => true,
        _ => false
    }).collect();

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

#[test]
pub fn expect_err_with_nameless_window() {
    let s = "---
windows:
  - : ls
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let result = main(&yaml, &"muxed".to_string(), false);
    assert!(result.is_err())
}

#[test]
pub fn expect_ok_with_empty_panes_syscommands() {
    let s = "---
windows:
  - cargo:
      layout: 'main-vertical'
      panes:
        -
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let result = main(&yaml, &"muxed".to_string(), false);
    assert!(result.is_ok())
}

#[test]
pub fn expect_no_send_keys_with_empty_panes_syscommands() {
    let s = "---
windows:
  - editor:
      panes:
        -
";

    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::SendKeys(_) => true,
        _ => false
    }).collect();

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

#[test]
pub fn expect_1_split_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::Split(_) => true,
        _ => false
    }).collect();

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

#[test]
pub fn expect_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 expect_1_session_with_panes_array() {
    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)
}

#[test]
pub fn expect_no_layout() {
    let s = "---
windows:
  - editor:
      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(), 0)
}

#[test]
pub fn expect_three_send_keys_commands_from_pre() {
    let s = "---
pre: 'ls'
windows:
  - editor:
      panes:
        -
        -
  - logs:
";
    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::SendKeys(_) => true,
        _ => false
    }).collect();

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

#[test]
pub fn expect_no_send_keys_with_blank_panes() {
    let s = "---
windows:
  - editor:
      panes: ['','','']
";

    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::SendKeys(_) => true,
        _ => false
    }).collect();

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

#[test]
pub fn expect_no_send_keys_with_blank_window() {
    let s = "---
windows:
  - editor: ''
";

    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::SendKeys(_) => true,
        _ => false
    }).collect();

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

#[test]
pub fn expect_full_directory_name() {
    let s = "---
root: ~/JustPlainSimple Technologies Inc./financials/ledgers
windows:
  - dir: ''
";

    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::SendKeys(_) => true,
        _ => false
    }).collect();

    let root = match remains[0] {
        Command::SendKeys(ref k) => k,
        _ => panic!("nope")
    };

    assert_eq!(root.exec, "cd \"~/JustPlainSimple Technologies Inc./financials/ledgers\"")
}
<|endoftext|>