<commit_before>use na::{Mat2, DMat};
use std::collections::HashMap;
use std::num::*;
use itertools::Itertools;
use std::cell::RefCell;
use std::rc::Rc;

mod test;

impl Default for DStreamProps {
    fn default() -> DStreamProps {
        DStreamProps {
            c_m: 3.0,
            c_l: 0.8,
            lambda: 0.998,
            beta: 0.3,
            i_bins: 10 as usize,
            j_bins: 10 as usize,
            i_range: (-10.0, 10.0),
            j_range: (-10.0, 10.0),
        }
    }
}
pub struct DStreamProps {
    c_m: f64,
    c_l: f64,
    lambda: f64,
    beta: f64,
    i_bins: usize,
    j_bins: usize,
    i_range: (f64, f64),
    j_range: (f64, f64),
}

#[derive(Clone)]
pub struct RawData {
    x: f64,
    y: f64,
    v: f64,
}

#[derive(Clone)]
#[derive(Debug)]
#[derive(Copy)]
struct GridPoint {
    t: u32,
    v: f64,
}
#[derive(Clone)]
struct GridData {
    i: usize,
    j: usize,
    v: f64,
}
struct TheWorld {
    g_vec: Vec<((usize, usize), DG)>,
}

#[test]
fn test_new_put_works() {
    let raw_data_1 = RawData{x: 1.0, y: -1.0, v: 123.45};
    let raw_data_2 = RawData{x: 5.0, y: -7.0, v: 23.45};
    let raw_data_3 = RawData{x: 5.0, y: -7.0, v: 83.45};
    let raw_data_4 = RawData{x: 15.0, y: -17.0, v: 83.45};

    let rd_vec = vec!(raw_data_1, raw_data_2, raw_data_3, raw_data_4);

    let t = 1;

    let default_vec : Vec<GridPoint> = Vec::new();
    let mut world = TheWorld{g_vec: Vec::new()};
    world.init(default_vec);
    let res1 = world.put(t, rd_vec.clone());
    let res2 = world.put(t + 1, rd_vec.clone());
    let res3 = world.put(t + 2, rd_vec.clone());

}

impl TheWorld {

    fn are_neighbors(dg1: DG, dg2: DG) -> bool {
        if (dg1.i == dg2.i) && ((dg1.j - dg2.j) as i32).abs() <= 1 {
            return true
        } else if (dg1.j == dg2.j) && ((dg1.i - dg2.i) as i32).abs() <= 1 {
            return true
        }
        false
    }

    #[test]
    fn test_compute_grid_indxs() {
        let loc = (-6., 6.);
        let i_rn = (-10.0, 10.0);
        let j_rn = (-10.0, 10.0);
        let i_bins = 10 as usize;
        let j_bins = 10 as usize;

        let result = TheWorld::compute_grid_indxs(loc, i_rn, j_rn, i_bins, j_bins).unwrap();
        assert_eq!((2, 8), result);

    }
    fn compute_grid_indxs(val: (f64, f64), i_range: (f64, f64), j_range: (f64, f64), i_bins: usize, j_bins: usize) -> Result<(usize, usize), String> {
        let i_size = (i_range.1 - i_range.0) / i_bins as f64;
        let i_number_of_sizes = ((val.0 - i_range.0) / i_size) as usize;
        let j_size = (j_range.1 - j_range.0) / j_bins as f64;
        let j_number_of_sizes = ((val.1 - j_range.0) / j_size) as usize;

        assert!(i_number_of_sizes < i_bins);
        assert!(j_number_of_sizes < j_bins);

        Ok((i_number_of_sizes, j_number_of_sizes))
    }

    fn init(&mut self, def_bucket: Vec<GridPoint>) {
        println!("init");
        let props: DStreamProps = DStreamProps { ..Default::default() };

        for i in 0..props.i_bins {
            for j in 0..props.j_bins {
                let z_clone = def_bucket.clone();
                let some_default_dg = DG {i: i, j: j,
                    updates_and_vals: z_clone, removed_as_spore_adic: Vec::new(),};
                (self.g_vec).push(((i as usize, j as usize), some_default_dg));
            }
        }
    }
    fn do_time_steps() {}
    fn do_one_time_step(t: u32, data: Vec<RawData>) {}

    fn get_by_idx(&mut self, idx: (usize, usize))-> DG {
        let  vec_f: Vec<((usize, usize), DG)> = self.g_vec.clone().into_iter().filter(|i| i.0 == idx).collect();
        vec_f[0].1.clone()
    }
    fn update_by_idx(&mut self, idx: (usize, usize), dg: DG) -> Result<(), String> {
        self.g_vec.retain(|i| i.0 != idx);
        self.g_vec.push((idx, dg));
        Ok(())
    }
    fn put(&mut self, t: u32, dat: Vec<RawData>) -> Result<(), String> {

        fn validate_range(loc2d: (f64, f64)) -> bool {

            let props: DStreamProps = DStreamProps { ..Default::default() };

            if ((loc2d.0 <= props.i_range.1) &&
            (loc2d.0 >= props.i_range.0) &&
            (loc2d.1 <= props.j_range.1) &&
            (loc2d.1 >= props.j_range.0)) {
                println!("valid!");
            } else {
                println!("invalid! -- bad input range");
                return false
            }
            return true
        }
        let with_idxs: Vec<GridData> = dat
            .iter()
            .filter(|rd| validate_range((rd.x, rd.y)))
            .map(|rd| self.which_idxs(rd).unwrap())
            .collect();

        let props: DStreamProps = DStreamProps { ..Default::default() };

        for (key, group) in with_idxs.iter().group_by(|gd| (gd.i, gd.j)) {
            println!("--put: key, g size: ({},{}) : {}", key.0, key.1, group.len());

            let the_vec_of_vals: Vec<f64> = group.iter().map(|t| t.v).collect();

            let mut some_default_dg = DG {i: key.0, j: key.1, updates_and_vals: Vec::<GridPoint>::new(), removed_as_spore_adic: Vec::new(),};

            let teh_dg: &mut DG = &mut self.get_by_idx(key);
            let update_dg_result = teh_dg.update(t, the_vec_of_vals);
            let udpate_world_result = self.update_by_idx(key, teh_dg.clone());
        }
        Ok(())
    }
    fn which_idxs(&self, dat: &RawData) -> Result<GridData, String> {

        let props: DStreamProps = DStreamProps { ..Default::default() };

        let idxs = TheWorld::compute_grid_indxs((dat.x, dat.y), props.i_range, props.j_range,props.i_bins, props.j_bins).unwrap();

        Ok((GridData{i:idxs.0,j:idxs.1,v:dat.v}))

    }
}

#[test]
fn test_dg_update_and_get() {
    let mut dg = DG {i: 0, j:0, updates_and_vals: Vec::new(), removed_as_spore_adic: Vec::new(),};
    dg.update(1, vec!(100.0));

    dg.update(10, vec!(200.0));
    dg.update(20, vec!(300.0));

    assert_eq!((1, 100.0), dg.get_last_update_and_value_to(2));
    assert_eq!((10, 200.0), dg.get_last_update_and_value_to(14));
    assert_eq!((20, 300.0), dg.get_last_update_and_value_to(22));

    let v_at_t = dg.get_at_time(30);
    println!("v: {}", v_at_t);
}

#[test]
fn test_removed_as_sporadic() {
    let mut dg = DG {i: 0, j:0, updates_and_vals: Vec::new(), removed_as_spore_adic: vec!(1, 2, 3, 4),};
    assert_eq!(1, dg.get_last_time_removed_as_sporadic_to(2));
    assert_eq!(4, dg.get_last_time_removed_as_sporadic_to(6));
}

#[derive(Debug)]
#[derive(Clone)]
//#[derive(Copy)] TODO: how to make this work?
struct DG {
    i: usize,
    j: usize,
    updates_and_vals: Vec<GridPoint>,
    removed_as_spore_adic: Vec<u32>,
}
#[derive(Debug)]
#[derive(Clone)]
enum GridLabel { Dense, Sparse, Transitional, }
#[derive(Debug)]
#[derive(Clone)]
enum GridStatus { Sporadic, Normal, }
impl DG {

    fn get_grid_label_at_time(&self, t:u32) -> GridLabel {
        let props: DStreamProps = DStreamProps { ..Default::default() };
        let n_size = (self.i * self.j) as f64;
        let d_m = props.c_m / (n_size * (1.0 - props.lambda));
        let d_l = props.c_l / (n_size * (1.0 - props.lambda));

        let d_t = self.get_at_time(t);

        if d_t >= d_m {
            GridLabel::Dense
        } else if d_t <= d_l {
            GridLabel::Sparse
        } else {
            assert!(d_t > d_l && d_t < d_t);
            GridLabel::Transitional
        }
    }

    fn get_at_time(&self, t: u32) -> f64 {
        let last_update_time_and_value = self.get_last_update_and_value_to(t);
        let coeff = self.coeff(t, last_update_time_and_value.0);
        coeff * last_update_time_and_value.1 + 1.0
    }
    fn update(&mut self, t: u32, vals: Vec<f64>) -> Result<(), String> {
        let sum = vals.clone().iter().fold(0.0, |sum, x| sum + x);
        self.updates_and_vals.push(GridPoint {t: t, v: sum});
        println!("  times updated dg: {}", self.updates_and_vals.len());
        Ok(())
    }

    fn get_last_update_and_value_to(&self, t: u32) -> (u32, f64) {
        let a: GridPoint = self.updates_and_vals.clone().into_iter().filter(|bp| bp.t < t).max_by_key(|bp| bp.t).unwrap();
        let t_l = a.t;
        let v_l = a.v;
        println!("  last update relative to {} is (t: {}, v: {})", t, t_l, v_l);
        (t_l, v_l)
    }

    fn get_last_time_removed_as_sporadic_to(&self, t: u32) -> u32 {
        let t = self.removed_as_spore_adic.clone().into_iter().filter(|&the_t| the_t < t).max().unwrap();
        t
    }

    fn coeff(&self, t_n: u32, t_l: u32) -> f64 {
        let props: DStreamProps = DStreamProps { ..Default::default() };
        props.lambda.powf((t_n - t_l) as f64)
    }

 }

#[test]
fn test_initialize_clustering() {

    let result = initialize_clustering(Vec::new());
    assert_eq!(Ok(()), result);
}

pub fn initialize_clustering(grid_list: Vec<RawData>) -> Result<(), String> {
    //update density of all grids in grid_list
    //assign each dense grid to a distinct cluster
    //label all other grids as NO_CLASS; bad grids!
    /*
    do until no change in cluster labels
          foreach cluster c
            foreach outside grid g of c
                foreach neighboring grid h of g
                    if h belongs to cluster c'
                        if |c| > |c'| label all grids in c' as c
                        else label all grids in c as c'
                    else if (h is translational) label h as in c
    end do
    */
    Ok(())
}

pub fn adjust_clustering(grid_list: Vec<RawData>) -> Result<(), String> {
    //update the density of all grids in grid_list
    /*
    foreach grid g whose attribute (dense/sparse/transitional) is changed since last call to adjust_clustering()
        if g is a sparse grid
            delete g from its cluster c, label g as NO_CLASS
            if (c becomes unconnected) split c into two clusters
        else if g is a dense grid
            among all neighboring grids of g, find out the grid h whoe cluster c_h has the largest size
            if h is a dense grid
                if (g is labelled as NO_CLASS) label g as in c_h
                else if (g is in cluster c and |c| > |c_h|)
                    label all grids in c_h as in c
                else if (g is in cluster c and |c| <= |c_h|)
                    label all grids in c as in c_h
            else if h is a transitional grid
                if ((g is NO_CLASS) and (h is an outside grid if g is added to c_h)) label g as in c_h
                else if (g is in cluster c and |c| >= |c_h|)
                    move h from cluster c_h to c
        else if (g is transitional grid)
            among neighboring clusters of g, find the largest one c'
            satisfying that g is an outside grid if added to it and
            label g as in c'
    */

    Ok(())
}<commit_msg>test are neighbors<commit_after>use na::{Mat2, DMat};
use std::collections::HashMap;
use std::num::*;
use itertools::Itertools;
use std::cell::RefCell;
use std::rc::Rc;

mod test;

impl Default for DStreamProps {
    fn default() -> DStreamProps {
        DStreamProps {
            c_m: 3.0,
            c_l: 0.8,
            lambda: 0.998,
            beta: 0.3,
            i_bins: 10 as usize,
            j_bins: 10 as usize,
            i_range: (-10.0, 10.0),
            j_range: (-10.0, 10.0),
        }
    }
}
pub struct DStreamProps {
    c_m: f64,
    c_l: f64,
    lambda: f64,
    beta: f64,
    i_bins: usize,
    j_bins: usize,
    i_range: (f64, f64),
    j_range: (f64, f64),
}

#[derive(Clone)]
pub struct RawData {
    x: f64,
    y: f64,
    v: f64,
}

#[derive(Clone)]
#[derive(Debug)]
#[derive(Copy)]
struct GridPoint {
    t: u32,
    v: f64,
}
#[derive(Clone)]
struct GridData {
    i: usize,
    j: usize,
    v: f64,
}
struct TheWorld {
    g_vec: Vec<((usize, usize), DG)>,
}

#[test]
fn test_new_put_works() {
    let raw_data_1 = RawData{x: 1.0, y: -1.0, v: 123.45};
    let raw_data_2 = RawData{x: 5.0, y: -7.0, v: 23.45};
    let raw_data_3 = RawData{x: 5.0, y: -7.0, v: 83.45};
    let raw_data_4 = RawData{x: 15.0, y: -17.0, v: 83.45};

    let rd_vec = vec!(raw_data_1, raw_data_2, raw_data_3, raw_data_4);

    let t = 1;

    let default_vec : Vec<GridPoint> = Vec::new();
    let mut world = TheWorld{g_vec: Vec::new()};
    world.init(default_vec);
    let res1 = world.put(t, rd_vec.clone());
    let res2 = world.put(t + 1, rd_vec.clone());
    let res3 = world.put(t + 2, rd_vec.clone());

}

#[test]
fn test_compute_grid_indxs() {
    let loc = (-6., 6.);
    let i_rn = (-10.0, 10.0);
    let j_rn = (-10.0, 10.0);
    let i_bins = 10 as usize;
    let j_bins = 10 as usize;

    let result = TheWorld::compute_grid_indxs(loc, i_rn, j_rn, i_bins, j_bins).unwrap();
    assert_eq!((2, 8), result);

}

#[test]
fn test_are_neighbors() {
    let dg1 = &DG {i: 0, j: 0,
        updates_and_vals: Vec::new(), removed_as_spore_adic: Vec::new(),};
    let dg2 = &DG {i: 1, j: 0,
        updates_and_vals: Vec::new(), removed_as_spore_adic: Vec::new(),};
    let dg3 = &DG {i: 1, j: 1,
        updates_and_vals: Vec::new(), removed_as_spore_adic: Vec::new(),};

    assert_eq!(true, TheWorld::are_neighbors(dg1, dg2));
    assert_eq!(true, TheWorld::are_neighbors(dg2, dg3));
    assert_eq!(false, TheWorld::are_neighbors(dg1, dg3));
}

impl TheWorld {

    fn are_neighbors(dg1: &DG, dg2: &DG) -> bool {
        if (dg1.i == dg2.i) && (dg1.j as i32 - dg2.j as i32).abs() <= 1 {
            return true
        } else if (dg1.j == dg2.j) && (dg1.i as i32 - dg2.i as i32).abs() <= 1 {
            return true
        }
        false
    }

    fn compute_grid_indxs(val: (f64, f64), i_range: (f64, f64), j_range: (f64, f64), i_bins: usize, j_bins: usize) -> Result<(usize, usize), String> {
        let i_size = (i_range.1 - i_range.0) / i_bins as f64;
        let i_number_of_sizes = ((val.0 - i_range.0) / i_size) as usize;
        let j_size = (j_range.1 - j_range.0) / j_bins as f64;
        let j_number_of_sizes = ((val.1 - j_range.0) / j_size) as usize;

        assert!(i_number_of_sizes < i_bins);
        assert!(j_number_of_sizes < j_bins);

        Ok((i_number_of_sizes, j_number_of_sizes))
    }

    fn init(&mut self, def_bucket: Vec<GridPoint>) {
        println!("init");
        let props: DStreamProps = DStreamProps { ..Default::default() };

        for i in 0..props.i_bins {
            for j in 0..props.j_bins {
                let z_clone = def_bucket.clone();
                let some_default_dg = DG {i: i, j: j,
                    updates_and_vals: z_clone, removed_as_spore_adic: Vec::new(),};
                (self.g_vec).push(((i as usize, j as usize), some_default_dg));
            }
        }
    }
    fn do_time_steps() {}
    fn do_one_time_step(t: u32, data: Vec<RawData>) {}

    fn get_by_idx(&mut self, idx: (usize, usize))-> DG {
        let  vec_f: Vec<((usize, usize), DG)> = self.g_vec.clone().into_iter().filter(|i| i.0 == idx).collect();
        vec_f[0].1.clone()
    }
    fn update_by_idx(&mut self, idx: (usize, usize), dg: DG) -> Result<(), String> {
        self.g_vec.retain(|i| i.0 != idx);
        self.g_vec.push((idx, dg));
        Ok(())
    }
    fn put(&mut self, t: u32, dat: Vec<RawData>) -> Result<(), String> {

        fn validate_range(loc2d: (f64, f64)) -> bool {

            let props: DStreamProps = DStreamProps { ..Default::default() };

            if ((loc2d.0 <= props.i_range.1) &&
            (loc2d.0 >= props.i_range.0) &&
            (loc2d.1 <= props.j_range.1) &&
            (loc2d.1 >= props.j_range.0)) {
                println!("valid!");
            } else {
                println!("invalid! -- bad input range");
                return false
            }
            return true
        }
        let with_idxs: Vec<GridData> = dat
            .iter()
            .filter(|rd| validate_range((rd.x, rd.y)))
            .map(|rd| self.which_idxs(rd).unwrap())
            .collect();

        let props: DStreamProps = DStreamProps { ..Default::default() };

        for (key, group) in with_idxs.iter().group_by(|gd| (gd.i, gd.j)) {
            println!("--put: key, g size: ({},{}) : {}", key.0, key.1, group.len());

            let the_vec_of_vals: Vec<f64> = group.iter().map(|t| t.v).collect();

            let mut some_default_dg = DG {i: key.0, j: key.1, updates_and_vals: Vec::<GridPoint>::new(), removed_as_spore_adic: Vec::new(),};

            let teh_dg: &mut DG = &mut self.get_by_idx(key);
            let update_dg_result = teh_dg.update(t, the_vec_of_vals);
            let udpate_world_result = self.update_by_idx(key, teh_dg.clone());
        }
        Ok(())
    }
    fn which_idxs(&self, dat: &RawData) -> Result<GridData, String> {

        let props: DStreamProps = DStreamProps { ..Default::default() };

        let idxs = TheWorld::compute_grid_indxs((dat.x, dat.y), props.i_range, props.j_range,props.i_bins, props.j_bins).unwrap();

        Ok((GridData{i:idxs.0,j:idxs.1,v:dat.v}))

    }
}

#[test]
fn test_dg_update_and_get() {
    let mut dg = DG {i: 0, j:0, updates_and_vals: Vec::new(), removed_as_spore_adic: Vec::new(),};
    dg.update(1, vec!(100.0));

    dg.update(10, vec!(200.0));
    dg.update(20, vec!(300.0));

    assert_eq!((1, 100.0), dg.get_last_update_and_value_to(2));
    assert_eq!((10, 200.0), dg.get_last_update_and_value_to(14));
    assert_eq!((20, 300.0), dg.get_last_update_and_value_to(22));

    let v_at_t = dg.get_at_time(30);
    println!("v: {}", v_at_t);
}

#[test]
fn test_removed_as_sporadic() {
    let mut dg = DG {i: 0, j:0, updates_and_vals: Vec::new(), removed_as_spore_adic: vec!(1, 2, 3, 4),};
    assert_eq!(1, dg.get_last_time_removed_as_sporadic_to(2));
    assert_eq!(4, dg.get_last_time_removed_as_sporadic_to(6));
}

#[derive(Debug)]
#[derive(Clone)]
//#[derive(Copy)] TODO: how to make this work?
struct DG {
    i: usize,
    j: usize,
    updates_and_vals: Vec<GridPoint>,
    removed_as_spore_adic: Vec<u32>,
}
#[derive(Debug)]
#[derive(Clone)]
enum GridLabel { Dense, Sparse, Transitional, }
#[derive(Debug)]
#[derive(Clone)]
enum GridStatus { Sporadic, Normal, }
impl DG {

    fn get_grid_label_at_time(&self, t:u32) -> GridLabel {
        let props: DStreamProps = DStreamProps { ..Default::default() };
        let n_size = (self.i * self.j) as f64;
        let d_m = props.c_m / (n_size * (1.0 - props.lambda));
        let d_l = props.c_l / (n_size * (1.0 - props.lambda));

        let d_t = self.get_at_time(t);

        if d_t >= d_m {
            GridLabel::Dense
        } else if d_t <= d_l {
            GridLabel::Sparse
        } else {
            assert!(d_t > d_l && d_t < d_t);
            GridLabel::Transitional
        }
    }

    fn get_at_time(&self, t: u32) -> f64 {
        let last_update_time_and_value = self.get_last_update_and_value_to(t);
        let coeff = self.coeff(t, last_update_time_and_value.0);
        coeff * last_update_time_and_value.1 + 1.0
    }
    fn update(&mut self, t: u32, vals: Vec<f64>) -> Result<(), String> {
        let sum = vals.clone().iter().fold(0.0, |sum, x| sum + x);
        self.updates_and_vals.push(GridPoint {t: t, v: sum});
        println!("  times updated dg: {}", self.updates_and_vals.len());
        Ok(())
    }

    fn get_last_update_and_value_to(&self, t: u32) -> (u32, f64) {
        let a: GridPoint = self.updates_and_vals.clone().into_iter().filter(|bp| bp.t < t).max_by_key(|bp| bp.t).unwrap();
        let t_l = a.t;
        let v_l = a.v;
        println!("  last update relative to {} is (t: {}, v: {})", t, t_l, v_l);
        (t_l, v_l)
    }

    fn get_last_time_removed_as_sporadic_to(&self, t: u32) -> u32 {
        let t = self.removed_as_spore_adic.clone().into_iter().filter(|&the_t| the_t < t).max().unwrap();
        t
    }

    fn coeff(&self, t_n: u32, t_l: u32) -> f64 {
        let props: DStreamProps = DStreamProps { ..Default::default() };
        props.lambda.powf((t_n - t_l) as f64)
    }

 }

#[test]
fn test_initialize_clustering() {

    let result = initialize_clustering(Vec::new());
    assert_eq!(Ok(()), result);
}

pub fn initialize_clustering(grid_list: Vec<RawData>) -> Result<(), String> {
    //update density of all grids in grid_list
    //assign each dense grid to a distinct cluster
    //label all other grids as NO_CLASS; bad grids!
    /*
    do until no change in cluster labels
          foreach cluster c
            foreach outside grid g of c
                foreach neighboring grid h of g
                    if h belongs to cluster c'
                        if |c| > |c'| label all grids in c' as c
                        else label all grids in c as c'
                    else if (h is translational) label h as in c
    end do
    */
    Ok(())
}

pub fn adjust_clustering(grid_list: Vec<RawData>) -> Result<(), String> {
    //update the density of all grids in grid_list
    /*
    foreach grid g whose attribute (dense/sparse/transitional) is changed since last call to adjust_clustering()
        if g is a sparse grid
            delete g from its cluster c, label g as NO_CLASS
            if (c becomes unconnected) split c into two clusters
        else if g is a dense grid
            among all neighboring grids of g, find out the grid h whoe cluster c_h has the largest size
            if h is a dense grid
                if (g is labelled as NO_CLASS) label g as in c_h
                else if (g is in cluster c and |c| > |c_h|)
                    label all grids in c_h as in c
                else if (g is in cluster c and |c| <= |c_h|)
                    label all grids in c as in c_h
            else if h is a transitional grid
                if ((g is NO_CLASS) and (h is an outside grid if g is added to c_h)) label g as in c_h
                else if (g is in cluster c and |c| >= |c_h|)
                    move h from cluster c_h to c
        else if (g is transitional grid)
            among neighboring clusters of g, find the largest one c'
            satisfying that g is an outside grid if added to it and
            label g as in c'
    */

    Ok(())
}<|endoftext|>


<commit_before>use clap::{App, Arg, ArgMatches};
use rav1e::*;
use std::fmt;
use std::fs::File;
use std::io;
use std::io::prelude::*;
use std::slice;
use std::sync::Arc;
use std::time::Instant;
use y4m;

pub struct EncoderIO {
  pub input: Box<dyn Read>,
  pub output: Box<dyn Write>,
  pub rec: Option<Box<dyn Write>>
}

pub struct CliOptions {
  pub io: EncoderIO,
  pub enc: EncoderConfig,
  pub limit: usize,
  pub verbose: bool,
}

pub fn parse_cli() -> CliOptions {
  let matches = App::new("rav1e")
    .version("0.1.0")
    .about("AV1 video encoder")
    .arg(
      Arg::with_name("INPUT")
        .help("Uncompressed YUV4MPEG2 video input")
        .required(true)
        .index(1)
    ).arg(
      Arg::with_name("OUTPUT")
        .help("Compressed AV1 in IVF video output")
        .short("o")
        .long("output")
        .required(true)
        .takes_value(true)
    ).arg(Arg::with_name("RECONSTRUCTION").short("r").takes_value(true))
    .arg(
      Arg::with_name("LIMIT")
        .help("Maximum number of frames to encode")
        .short("l")
        .long("limit")
        .takes_value(true)
        .default_value("0")
    ).arg(
      Arg::with_name("QP")
        .help("Quantizer (0-255)")
        .long("quantizer")
        .takes_value(true)
        .default_value("100")
    ).arg(
      Arg::with_name("SPEED")
        .help("Speed level (0(slow)-10(fast))")
        .short("s")
        .long("speed")
        .takes_value(true)
        .default_value("3")
    ).arg(
      Arg::with_name("KEYFRAME_INTERVAL")
        .help("Keyframe interval")
        .short("I")
        .long("keyint")
        .takes_value(true)
        .default_value("30")
    ).arg(
      Arg::with_name("LOW_LATENCY")
        .help("low latency mode. true or false")
        .long("low_latency")
        .takes_value(true)
        .default_value("true")
    ).arg(
      Arg::with_name("TUNE")
        .help("Quality tuning (Will enforce partition sizes >= 8x8)")
        .long("tune")
        .possible_values(&Tune::variants())
        .default_value("psnr")
        .case_insensitive(true)
    ).arg(
      Arg::with_name("VERBOSE")
        .help("verbose logging, output info for every frame")
        .long("verbose")
        .short("v")
    ).get_matches();

  let io = EncoderIO {
    input: match matches.value_of("INPUT").unwrap() {
      "-" => Box::new(io::stdin()) as Box<dyn Read>,
      f => Box::new(File::open(&f).unwrap()) as Box<dyn Read>
    },
    output: match matches.value_of("OUTPUT").unwrap() {
      "-" => Box::new(io::stdout()) as Box<dyn Write>,
      f => Box::new(File::create(&f).unwrap()) as Box<dyn Write>
    },
    rec: matches
      .value_of("RECONSTRUCTION")
      .map(|f| Box::new(File::create(&f).unwrap()) as Box<dyn Write>)
  };

  CliOptions {
    io,
    enc: parse_config(&matches),
    limit: matches.value_of("LIMIT").unwrap().parse().unwrap(),
    verbose: matches.is_present("VERBOSE"),
  }
}

fn parse_config(matches: &ArgMatches) -> EncoderConfig {
  let speed = matches.value_of("SPEED").unwrap().parse().unwrap();
  let quantizer = matches.value_of("QP").unwrap().parse().unwrap();

  // Validate arguments
  if quantizer == 0 {
    unimplemented!("Lossless encoding not yet implemented");
  } else if quantizer > 255 || speed > 10 {
    panic!("argument out of range");
  }

  let mut cfg = EncoderConfig::with_speed_preset(speed);
  cfg.key_frame_interval = matches.value_of("KEYFRAME_INTERVAL").unwrap().parse().unwrap();
  cfg.low_latency = matches.value_of("LOW_LATENCY").unwrap().parse().unwrap();
  cfg.tune = matches.value_of("TUNE").unwrap().parse().unwrap();
  cfg.quantizer = quantizer;
  cfg
}

#[derive(Debug, Clone, Copy)]
pub struct FrameSummary {
  /// Frame size in bytes
  pub size: usize,
  pub number: u64,
  pub frame_type: FrameType
}

impl From<Packet> for FrameSummary {
  fn from(packet: Packet) -> Self {
    Self {
      size: packet.data.len(),
      number: packet.number,
      frame_type: packet.frame_type,
    }
  }
}

impl fmt::Display for FrameSummary {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    write!(
      f,
      "Frame {} - {} - {} bytes",
      self.number,
      self.frame_type,
      self.size
    )
  }
}

/// Encode and write a frame.
/// Returns frame information in a `Result`.
pub fn process_frame(
  ctx: &mut Context, output_file: &mut dyn Write,
  y4m_dec: &mut y4m::Decoder<'_, Box<dyn Read>>,
  mut y4m_enc: Option<&mut y4m::Encoder<'_, Box<dyn Write>>>
) -> Result<Vec<FrameSummary>, ()> {
  let width = y4m_dec.get_width();
  let height = y4m_dec.get_height();
  let y4m_bits = y4m_dec.get_bit_depth();
  let y4m_bytes = y4m_dec.get_bytes_per_sample();
  let bit_depth = y4m_dec.get_colorspace().get_bit_depth();

  if ctx.needs_more_frames(ctx.get_frame_count()) {
    match y4m_dec.read_frame() {
      Ok(y4m_frame) => {
        let y4m_y = y4m_frame.get_y_plane();
        let y4m_u = y4m_frame.get_u_plane();
        let y4m_v = y4m_frame.get_v_plane();
        let mut input = ctx.new_frame();
        {
          let input = Arc::get_mut(&mut input).unwrap();
          input.planes[0].copy_from_raw_u8(&y4m_y, width * y4m_bytes, y4m_bytes);
          input.planes[1].copy_from_raw_u8(
            &y4m_u,
            width * y4m_bytes / 2,
            y4m_bytes
          );
          input.planes[2].copy_from_raw_u8(
            &y4m_v,
            width * y4m_bytes / 2,
            y4m_bytes
          );
        }

        match y4m_bits {
          8 | 10 | 12 => {}
          _ => panic!("unknown input bit depth!")
        }

        let _ = ctx.send_frame(input);
        true
      }
      _ => {
        let frames_to_be_coded = ctx.get_frame_count();
        ctx.set_frames_to_be_coded(frames_to_be_coded);
        ctx.flush();
        false
      }
    }
  } else {
    ctx.flush();
    false
  };

  let mut has_data = true;
  let mut frame_summaries = Vec::new();
  while has_data {
    let pkt_wrapped = ctx.receive_packet();
    match pkt_wrapped {
      Ok(pkt) => {
        write_ivf_frame(output_file, pkt.number as u64, pkt.data.as_ref());
        if let Some(y4m_enc_uw) = y4m_enc.as_mut() {
          if let Some(ref rec) = pkt.rec {
            let pitch_y = if bit_depth > 8 { width * 2 } else { width };
            let pitch_uv = pitch_y / 2;

            let (mut rec_y, mut rec_u, mut rec_v) = (
              vec![128u8; pitch_y * height],
              vec![128u8; pitch_uv * (height / 2)],
              vec![128u8; pitch_uv * (height / 2)]
            );

            let (stride_y, stride_u, stride_v) = (
              rec.planes[0].cfg.stride,
              rec.planes[1].cfg.stride,
              rec.planes[2].cfg.stride
            );

            for (line, line_out) in rec.planes[0]
              .data_origin()
              .chunks(stride_y)
              .zip(rec_y.chunks_mut(pitch_y))
            {
              if bit_depth > 8 {
                unsafe {
                  line_out.copy_from_slice(slice::from_raw_parts::<u8>(
                    line.as_ptr() as (*const u8),
                    pitch_y
                  ));
                }
              } else {
                line_out.copy_from_slice(
                  &line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_y]
                );
              }
            }
            for (line, line_out) in rec.planes[1]
              .data_origin()
              .chunks(stride_u)
              .zip(rec_u.chunks_mut(pitch_uv))
            {
              if bit_depth > 8 {
                unsafe {
                  line_out.copy_from_slice(slice::from_raw_parts::<u8>(
                    line.as_ptr() as (*const u8),
                    pitch_uv
                  ));
                }
              } else {
                line_out.copy_from_slice(
                  &line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_uv]
                );
              }
            }
            for (line, line_out) in rec.planes[2]
              .data_origin()
              .chunks(stride_v)
              .zip(rec_v.chunks_mut(pitch_uv))
            {
              if bit_depth > 8 {
                unsafe {
                  line_out.copy_from_slice(slice::from_raw_parts::<u8>(
                    line.as_ptr() as (*const u8),
                    pitch_uv
                  ));
                }
              } else {
                line_out.copy_from_slice(
                  &line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_uv]
                );
              }
            }

            let rec_frame = y4m::Frame::new([&rec_y, &rec_u, &rec_v], None);
            y4m_enc_uw.write_frame(&rec_frame).unwrap();
          }
        }
        frame_summaries.push(pkt.into());
      },
      _ => { has_data = false; }
    }
  }
  Ok(frame_summaries)
}

#[derive(Debug, Clone)]
pub struct ProgressInfo {
  /// Frame rate of the video
  frame_rate: y4m::Ratio,
  /// The length of the whole video, in frames, if known
  total_frames: Option<usize>,
  /// The time the encode was started
  time_started: Instant,
  /// List of frames encoded so far
  frame_info: Vec<FrameSummary>,
  /// Video size so far in bytes.
  ///
  /// This value will be updated in the CLI very frequently, so we cache the previous value
  /// to reduce the overall complexity.
  encoded_size: usize,
}

impl ProgressInfo {
  pub fn new(frame_rate: y4m::Ratio, total_frames: Option<usize>) -> Self {
    Self {
      frame_rate,
      total_frames,
      time_started: Instant::now(),
      frame_info: Vec::with_capacity(total_frames.unwrap_or_default()),
      encoded_size: 0,
    }
  }

  pub fn add_frame(&mut self, frame: FrameSummary) {
    self.encoded_size += frame.size;
    self.frame_info.push(frame);
  }

  pub fn frames_encoded(&self) -> usize {
    self.frame_info.len()
  }

  pub fn encoding_fps(&self) -> f64 {
    let seconds = Instant::now().duration_since(self.time_started).as_secs();
    self.frame_info.len() as f64 / seconds as f64
  }

  pub fn video_fps(&self) -> f64 {
    self.frame_rate.num as f64 / self.frame_rate.den as f64
  }

  /// Returns the bitrate of the frames so far, in bits/second
  pub fn bitrate(&self) -> usize {
    let bits = self.encoded_size * 8;
    let seconds = self.frame_info.len() as f64 / self.video_fps();
    (bits as f64 / seconds) as usize
  }

  /// Estimates the final filesize in bytes, if the number of frames is known
  pub fn estimated_size(&self) -> usize {
    self.total_frames
      .map(|frames| self.encoded_size * frames / self.frames_encoded())
      .unwrap_or_default()
  }

  /// Number of frames of given type which appear in the video
  pub fn get_frame_type_count(&self, frame_type: FrameType) -> usize {
    self.frame_info.iter()
      .filter(|frame| frame.frame_type == frame_type)
      .count()
  }

  /// Size in bytes of all frames of given frame type
  pub fn get_frame_type_size(&self, frame_type: FrameType) -> usize {
    self.frame_info.iter()
      .filter(|frame| frame.frame_type == frame_type)
      .map(|frame| frame.size)
      .sum()
  }

  pub fn print_stats(&self) -> String {
    let (key, key_size) = (self.get_frame_type_count(FrameType::KEY), self.get_frame_type_size(FrameType::KEY));
    let (inter, inter_size) = (self.get_frame_type_count(FrameType::INTER), self.get_frame_type_size(FrameType::INTER));
    let (ionly, ionly_size) = (self.get_frame_type_count(FrameType::INTRA_ONLY), self.get_frame_type_size(FrameType::INTRA_ONLY));
    let (switch, switch_size) = (self.get_frame_type_count(FrameType::SWITCH), self.get_frame_type_size(FrameType::SWITCH));
    format!("\
    Key Frames: {:>6}    avg size: {:>7} B\n\
    Inter:      {:>6}    avg size: {:>7} B\n\
    Intra Only: {:>6}    avg size: {:>7} B\n\
    Switch:     {:>6}    avg size: {:>7} B",
      key, key_size / key,
      inter, inter_size / inter,
      ionly, ionly_size / key,
      switch, switch_size / key
    )
  }
}

impl fmt::Display for ProgressInfo {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    if let Some(total_frames) = self.total_frames {
      write!(
        f,
        "encoded {}/{} frames, {:.2} fps, {:.2} Kb/s, est. size: {:.2} MB",
        self.frames_encoded(),
        total_frames,
        self.encoding_fps(),
        self.bitrate() as f64 / 1024f64,
        self.estimated_size() as f64 / (1024 * 1024) as f64
      )
    } else {
      write!(
        f,
        "encoded {} frames, {:.2} fps, {:.2} Kb/s",
        self.frames_encoded(),
        self.encoding_fps(),
        self.bitrate() as f64 / 1024f64
      )
    }
  }
}
<commit_msg>Fix division by zero in CLI output when no inter frame is coded<commit_after>use clap::{App, Arg, ArgMatches};
use rav1e::*;
use std::fmt;
use std::fs::File;
use std::io;
use std::io::prelude::*;
use std::slice;
use std::sync::Arc;
use std::time::Instant;
use y4m;

pub struct EncoderIO {
  pub input: Box<dyn Read>,
  pub output: Box<dyn Write>,
  pub rec: Option<Box<dyn Write>>
}

pub struct CliOptions {
  pub io: EncoderIO,
  pub enc: EncoderConfig,
  pub limit: usize,
  pub verbose: bool,
}

pub fn parse_cli() -> CliOptions {
  let matches = App::new("rav1e")
    .version("0.1.0")
    .about("AV1 video encoder")
    .arg(
      Arg::with_name("INPUT")
        .help("Uncompressed YUV4MPEG2 video input")
        .required(true)
        .index(1)
    ).arg(
      Arg::with_name("OUTPUT")
        .help("Compressed AV1 in IVF video output")
        .short("o")
        .long("output")
        .required(true)
        .takes_value(true)
    ).arg(Arg::with_name("RECONSTRUCTION").short("r").takes_value(true))
    .arg(
      Arg::with_name("LIMIT")
        .help("Maximum number of frames to encode")
        .short("l")
        .long("limit")
        .takes_value(true)
        .default_value("0")
    ).arg(
      Arg::with_name("QP")
        .help("Quantizer (0-255)")
        .long("quantizer")
        .takes_value(true)
        .default_value("100")
    ).arg(
      Arg::with_name("SPEED")
        .help("Speed level (0(slow)-10(fast))")
        .short("s")
        .long("speed")
        .takes_value(true)
        .default_value("3")
    ).arg(
      Arg::with_name("KEYFRAME_INTERVAL")
        .help("Keyframe interval")
        .short("I")
        .long("keyint")
        .takes_value(true)
        .default_value("30")
    ).arg(
      Arg::with_name("LOW_LATENCY")
        .help("low latency mode. true or false")
        .long("low_latency")
        .takes_value(true)
        .default_value("true")
    ).arg(
      Arg::with_name("TUNE")
        .help("Quality tuning (Will enforce partition sizes >= 8x8)")
        .long("tune")
        .possible_values(&Tune::variants())
        .default_value("psnr")
        .case_insensitive(true)
    ).arg(
      Arg::with_name("VERBOSE")
        .help("verbose logging, output info for every frame")
        .long("verbose")
        .short("v")
    ).get_matches();

  let io = EncoderIO {
    input: match matches.value_of("INPUT").unwrap() {
      "-" => Box::new(io::stdin()) as Box<dyn Read>,
      f => Box::new(File::open(&f).unwrap()) as Box<dyn Read>
    },
    output: match matches.value_of("OUTPUT").unwrap() {
      "-" => Box::new(io::stdout()) as Box<dyn Write>,
      f => Box::new(File::create(&f).unwrap()) as Box<dyn Write>
    },
    rec: matches
      .value_of("RECONSTRUCTION")
      .map(|f| Box::new(File::create(&f).unwrap()) as Box<dyn Write>)
  };

  CliOptions {
    io,
    enc: parse_config(&matches),
    limit: matches.value_of("LIMIT").unwrap().parse().unwrap(),
    verbose: matches.is_present("VERBOSE"),
  }
}

fn parse_config(matches: &ArgMatches) -> EncoderConfig {
  let speed = matches.value_of("SPEED").unwrap().parse().unwrap();
  let quantizer = matches.value_of("QP").unwrap().parse().unwrap();

  // Validate arguments
  if quantizer == 0 {
    unimplemented!("Lossless encoding not yet implemented");
  } else if quantizer > 255 || speed > 10 {
    panic!("argument out of range");
  }

  let mut cfg = EncoderConfig::with_speed_preset(speed);
  cfg.key_frame_interval = matches.value_of("KEYFRAME_INTERVAL").unwrap().parse().unwrap();
  cfg.low_latency = matches.value_of("LOW_LATENCY").unwrap().parse().unwrap();
  cfg.tune = matches.value_of("TUNE").unwrap().parse().unwrap();
  cfg.quantizer = quantizer;
  cfg
}

#[derive(Debug, Clone, Copy)]
pub struct FrameSummary {
  /// Frame size in bytes
  pub size: usize,
  pub number: u64,
  pub frame_type: FrameType
}

impl From<Packet> for FrameSummary {
  fn from(packet: Packet) -> Self {
    Self {
      size: packet.data.len(),
      number: packet.number,
      frame_type: packet.frame_type,
    }
  }
}

impl fmt::Display for FrameSummary {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    write!(
      f,
      "Frame {} - {} - {} bytes",
      self.number,
      self.frame_type,
      self.size
    )
  }
}

/// Encode and write a frame.
/// Returns frame information in a `Result`.
pub fn process_frame(
  ctx: &mut Context, output_file: &mut dyn Write,
  y4m_dec: &mut y4m::Decoder<'_, Box<dyn Read>>,
  mut y4m_enc: Option<&mut y4m::Encoder<'_, Box<dyn Write>>>
) -> Result<Vec<FrameSummary>, ()> {
  let width = y4m_dec.get_width();
  let height = y4m_dec.get_height();
  let y4m_bits = y4m_dec.get_bit_depth();
  let y4m_bytes = y4m_dec.get_bytes_per_sample();
  let bit_depth = y4m_dec.get_colorspace().get_bit_depth();

  if ctx.needs_more_frames(ctx.get_frame_count()) {
    match y4m_dec.read_frame() {
      Ok(y4m_frame) => {
        let y4m_y = y4m_frame.get_y_plane();
        let y4m_u = y4m_frame.get_u_plane();
        let y4m_v = y4m_frame.get_v_plane();
        let mut input = ctx.new_frame();
        {
          let input = Arc::get_mut(&mut input).unwrap();
          input.planes[0].copy_from_raw_u8(&y4m_y, width * y4m_bytes, y4m_bytes);
          input.planes[1].copy_from_raw_u8(
            &y4m_u,
            width * y4m_bytes / 2,
            y4m_bytes
          );
          input.planes[2].copy_from_raw_u8(
            &y4m_v,
            width * y4m_bytes / 2,
            y4m_bytes
          );
        }

        match y4m_bits {
          8 | 10 | 12 => {}
          _ => panic!("unknown input bit depth!")
        }

        let _ = ctx.send_frame(input);
        true
      }
      _ => {
        let frames_to_be_coded = ctx.get_frame_count();
        ctx.set_frames_to_be_coded(frames_to_be_coded);
        ctx.flush();
        false
      }
    }
  } else {
    ctx.flush();
    false
  };

  let mut has_data = true;
  let mut frame_summaries = Vec::new();
  while has_data {
    let pkt_wrapped = ctx.receive_packet();
    match pkt_wrapped {
      Ok(pkt) => {
        write_ivf_frame(output_file, pkt.number as u64, pkt.data.as_ref());
        if let Some(y4m_enc_uw) = y4m_enc.as_mut() {
          if let Some(ref rec) = pkt.rec {
            let pitch_y = if bit_depth > 8 { width * 2 } else { width };
            let pitch_uv = pitch_y / 2;

            let (mut rec_y, mut rec_u, mut rec_v) = (
              vec![128u8; pitch_y * height],
              vec![128u8; pitch_uv * (height / 2)],
              vec![128u8; pitch_uv * (height / 2)]
            );

            let (stride_y, stride_u, stride_v) = (
              rec.planes[0].cfg.stride,
              rec.planes[1].cfg.stride,
              rec.planes[2].cfg.stride
            );

            for (line, line_out) in rec.planes[0]
              .data_origin()
              .chunks(stride_y)
              .zip(rec_y.chunks_mut(pitch_y))
            {
              if bit_depth > 8 {
                unsafe {
                  line_out.copy_from_slice(slice::from_raw_parts::<u8>(
                    line.as_ptr() as (*const u8),
                    pitch_y
                  ));
                }
              } else {
                line_out.copy_from_slice(
                  &line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_y]
                );
              }
            }
            for (line, line_out) in rec.planes[1]
              .data_origin()
              .chunks(stride_u)
              .zip(rec_u.chunks_mut(pitch_uv))
            {
              if bit_depth > 8 {
                unsafe {
                  line_out.copy_from_slice(slice::from_raw_parts::<u8>(
                    line.as_ptr() as (*const u8),
                    pitch_uv
                  ));
                }
              } else {
                line_out.copy_from_slice(
                  &line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_uv]
                );
              }
            }
            for (line, line_out) in rec.planes[2]
              .data_origin()
              .chunks(stride_v)
              .zip(rec_v.chunks_mut(pitch_uv))
            {
              if bit_depth > 8 {
                unsafe {
                  line_out.copy_from_slice(slice::from_raw_parts::<u8>(
                    line.as_ptr() as (*const u8),
                    pitch_uv
                  ));
                }
              } else {
                line_out.copy_from_slice(
                  &line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_uv]
                );
              }
            }

            let rec_frame = y4m::Frame::new([&rec_y, &rec_u, &rec_v], None);
            y4m_enc_uw.write_frame(&rec_frame).unwrap();
          }
        }
        frame_summaries.push(pkt.into());
      },
      _ => { has_data = false; }
    }
  }
  Ok(frame_summaries)
}

#[derive(Debug, Clone)]
pub struct ProgressInfo {
  /// Frame rate of the video
  frame_rate: y4m::Ratio,
  /// The length of the whole video, in frames, if known
  total_frames: Option<usize>,
  /// The time the encode was started
  time_started: Instant,
  /// List of frames encoded so far
  frame_info: Vec<FrameSummary>,
  /// Video size so far in bytes.
  ///
  /// This value will be updated in the CLI very frequently, so we cache the previous value
  /// to reduce the overall complexity.
  encoded_size: usize,
}

impl ProgressInfo {
  pub fn new(frame_rate: y4m::Ratio, total_frames: Option<usize>) -> Self {
    Self {
      frame_rate,
      total_frames,
      time_started: Instant::now(),
      frame_info: Vec::with_capacity(total_frames.unwrap_or_default()),
      encoded_size: 0,
    }
  }

  pub fn add_frame(&mut self, frame: FrameSummary) {
    self.encoded_size += frame.size;
    self.frame_info.push(frame);
  }

  pub fn frames_encoded(&self) -> usize {
    self.frame_info.len()
  }

  pub fn encoding_fps(&self) -> f64 {
    let seconds = Instant::now().duration_since(self.time_started).as_secs();
    self.frame_info.len() as f64 / seconds as f64
  }

  pub fn video_fps(&self) -> f64 {
    self.frame_rate.num as f64 / self.frame_rate.den as f64
  }

  /// Returns the bitrate of the frames so far, in bits/second
  pub fn bitrate(&self) -> usize {
    let bits = self.encoded_size * 8;
    let seconds = self.frame_info.len() as f64 / self.video_fps();
    (bits as f64 / seconds) as usize
  }

  /// Estimates the final filesize in bytes, if the number of frames is known
  pub fn estimated_size(&self) -> usize {
    self.total_frames
      .map(|frames| self.encoded_size * frames / self.frames_encoded())
      .unwrap_or_default()
  }

  /// Number of frames of given type which appear in the video
  pub fn get_frame_type_count(&self, frame_type: FrameType) -> usize {
    self.frame_info.iter()
      .filter(|frame| frame.frame_type == frame_type)
      .count()
  }

  /// Size in bytes of all frames of given frame type
  pub fn get_frame_type_size(&self, frame_type: FrameType) -> usize {
    self.frame_info.iter()
      .filter(|frame| frame.frame_type == frame_type)
      .map(|frame| frame.size)
      .sum()
  }

  pub fn print_stats(&self) -> String {
    let (key, key_size) = (self.get_frame_type_count(FrameType::KEY), self.get_frame_type_size(FrameType::KEY));
    let (inter, inter_size) = (self.get_frame_type_count(FrameType::INTER), self.get_frame_type_size(FrameType::INTER));
    let (ionly, ionly_size) = (self.get_frame_type_count(FrameType::INTRA_ONLY), self.get_frame_type_size(FrameType::INTRA_ONLY));
    let (switch, switch_size) = (self.get_frame_type_count(FrameType::SWITCH), self.get_frame_type_size(FrameType::SWITCH));
    format!("\
    Key Frames: {:>6}    avg size: {:>7} B\n\
    Inter:      {:>6}    avg size: {:>7} B\n\
    Intra Only: {:>6}    avg size: {:>7} B\n\
    Switch:     {:>6}    avg size: {:>7} B",
      key, key_size / key,
      inter, inter_size.checked_div(inter).unwrap_or(0),
      ionly, ionly_size / key,
      switch, switch_size / key
    )
  }
}

impl fmt::Display for ProgressInfo {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    if let Some(total_frames) = self.total_frames {
      write!(
        f,
        "encoded {}/{} frames, {:.2} fps, {:.2} Kb/s, est. size: {:.2} MB",
        self.frames_encoded(),
        total_frames,
        self.encoding_fps(),
        self.bitrate() as f64 / 1024f64,
        self.estimated_size() as f64 / (1024 * 1024) as f64
      )
    } else {
      write!(
        f,
        "encoded {} frames, {:.2} fps, {:.2} Kb/s",
        self.frames_encoded(),
        self.encoding_fps(),
        self.bitrate() as f64 / 1024f64
      )
    }
  }
}
<|endoftext|>


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

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

extern crate netflowv9;
use netflowv9::*;

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

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

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

type RawDataTemplateIter = Vec<u8>;

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

type Output = BufWriter<ZlibEncoder<File>>;

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

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

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

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

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

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

        self.template_ids_seen.clear();

        Ok(())
    }

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

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

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

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

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

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

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

        Ok(buf.len())
    }

    fn flush(&mut self) -> std::io::Result<()>{
        self.output.flush()
    }
}
<commit_msg>Bugfix. Use the correct length.<commit_after>use std::collections::{HashMap, HashSet};
use std::fmt::Display;
use std::fs::{File, OpenOptions};
use std::io::prelude::*;
use std::io::{BufWriter, SeekFrom};
use std::mem;
use std::net::{UdpSocket, IpAddr};
use std::time::{Duration, Instant};

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

extern crate netflowv9;
use netflowv9::*;

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

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

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

type RawDataTemplateIter = Vec<u8>;

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

type Output = BufWriter<ZlibEncoder<File>>;

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

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

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

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

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

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

        self.template_ids_seen.clear();

        Ok(())
    }

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

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

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

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

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

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

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

        Ok(buf.len())
    }

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


<commit_before>// xfail-pretty
// xfail-test

use std::cast::transmute;
use std::libc::{STDOUT_FILENO, c_int, fdopen, fgets, fopen, fputc, fwrite};
use std::libc::{size_t};
use std::ptr::null;

static LINE_LEN: u32 = 80;

static COMPLEMENTS: [u8, ..256] = [
    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,

    0,
    'T' as u8,
    'V' as u8,
    'G' as u8,
    'H' as u8,
    0,
    0,
    'C' as u8,
    'D' as u8,
    0,
    0,
    'M' as u8,
    0,
    'K' as u8,
    'N' as u8,
    0,
    0,
    0,
    'Y' as u8,
    'S' as u8,
    'A' as u8,
    'A' as u8,
    'B' as u8,
    'W' as u8,
    0,
    'R' as u8,
    0,
    0,
    0,
    0,
    0,
    0,

    0,
    'T' as u8,
    'V' as u8,
    'G' as u8,
    'H' as u8,
    0,
    0,
    'C' as u8,
    'D' as u8,
    0,
    0,
    'M' as u8,
    0,
    'K' as u8,
    'N' as u8,
    0,
    0,
    0,
    'Y' as u8,
    'S' as u8,
    'A' as u8,
    'A' as u8,
    'B' as u8,
    'W' as u8,
    0,
    'R' as u8,
    0,
    0,
    0,
    0,
    0,
    0,

    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
];

fn main() {
    unsafe {
        let mode = "r";
        //let stdin = fdopen(STDIN_FILENO as c_int, transmute(&mode[0]));
        let path = "reversecomplement-input.txt";
        let stdin = fopen(transmute(&path[0]), transmute(&mode[0]));
        let mode = "w";
        let stdout = fdopen(STDOUT_FILENO as c_int, transmute(&mode[0]));

        let mut out: ~[u8] = ~[];
        out.reserve(12777888);
        let mut pos = 0;

        loop {
            let needed = pos + (LINE_LEN as uint) + 1;
            if out.capacity() < needed {
                out.reserve_at_least(needed);
            }

            let mut ptr = out.unsafe_mut_ref(pos);
            if fgets(transmute(ptr), LINE_LEN as c_int, stdin) == null() {
                break;
            }

            // Don't change lines that begin with '>' or ';'.
            let first = *ptr;
            if first == ('>' as u8) {
                while *ptr != 0 {
                    ptr = ptr.offset(1);
                }
                *ptr = '\n' as u8;

                pos = (ptr as uint) - (out.unsafe_ref(0) as uint);
                fwrite(transmute(out.unsafe_ref(0)),
                       1,
                       pos as size_t,
                       stdout);

                pos = 0;
                continue;
            }

            // Complement other lines.
            loop {
                let ch = *ptr;
                if ch == 0 {
                    break;
                }
                *ptr = COMPLEMENTS.unsafe_get(ch as uint);
                ptr = ptr.offset(1);
            }
            *ptr = '\n' as u8;

            pos = (ptr as uint) - (out.unsafe_ref(0) as uint);
        }

        fwrite(transmute(out.unsafe_ref(0)), 1, pos as size_t, stdout);
    }
}
<commit_msg>auto merge of #10799 : TeXitoi/rust/shootout-reverse-complement-resurected, r=alexcrichton<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::iter::range_step;
use std::io::{stdin, stdout, File};

static LINE_LEN: uint = 60;

fn make_complements() -> [u8, ..256] {
    let transforms = [
        ('A', 'T'), ('C', 'G'), ('G', 'C'), ('T', 'A'),
        ('U', 'A'), ('M', 'K'), ('R', 'Y'), ('W', 'W'),
        ('S', 'S'), ('Y', 'R'), ('K', 'M'), ('V', 'B'),
        ('H', 'D'), ('D', 'H'), ('B', 'V'), ('N', 'N'),
        ('\n', '\n')];
    let mut complements: [u8, ..256] = [0, ..256];
    for (i, c) in complements.mut_iter().enumerate() {
        *c = i as u8;
    }
    let lower = 'A' as u8 - 'a' as u8;
    for &(from, to) in transforms.iter() {
        complements[from as u8] = to as u8;
        complements[from as u8 - lower] = to as u8;
    }
    complements
}

fn main() {
    let complements = make_complements();
    let mut data = if std::os::getenv("RUST_BENCH").is_some() {
        File::open(&Path::init("shootout-k-nucleotide.data")).read_to_end()
    } else {
        stdin().read_to_end()
    };

    for seq in data.mut_split(|c| *c == '>' as u8) {
        // skip header and last \n
        let begin = match seq.iter().position(|c| *c == '\n' as u8) {
            None => continue,
            Some(c) => c
        };
        let len = seq.len();
        let seq = seq.mut_slice(begin + 1, len - 1);

        // arrange line breaks
        let len = seq.len();
        let off = LINE_LEN - len % (LINE_LEN + 1);
        for i in range_step(LINE_LEN, len, LINE_LEN + 1) {
            for j in std::iter::count(i, -1).take(off) {
                seq[j] = seq[j - 1];
            }
            seq[i - off] = '\n' as u8;
        }

        // reverse complement, as
        //    seq.reverse(); for c in seq.mut_iter() {*c = complements[*c]}
        // but faster:
        let mut it = seq.mut_iter();
        loop {
            match (it.next(), it.next_back()) {
                (Some(front), Some(back)) => {
                    let tmp = complements[*front];
                    *front = complements[*back];
                    *back = tmp;
                }
                _ => break // vector exhausted.
            }
        }
    }

    stdout().write(data);
}
<|endoftext|>


<commit_before>extern mod std;
use std::time::precise_time_ns;
use cast::{reinterpret_cast, transmute};
use vec::raw::buf_as_slice;
use ptr::to_unsafe_ptr;

pub struct Vector { x: float, y: float, z: float, w: float }

pub struct VecMatch { x: float, y: float, z: float, w: float }

pub impl VecMatch: Index<uint, float> {
    #[inline(always)]
    pure fn index(i: uint) -> float {
        match i {
            0 => self.x,
            1 => self.y,
            2 => self.z,
            3 => self.w,
            _ => fail(~"Vector index out of bounds.")
        }
    }
}

pub struct VecReinterpret { x: float, y: float, z: float, w: float }

pub impl VecReinterpret: Index<uint, float> {
    #[inline(always)]
    pure fn index(i: uint) -> float unsafe {
        (reinterpret_cast::<VecReinterpret, [float * 4]>(&self))[i]
    }
}

pub struct VecTransmute { x: float, y: float, z: float, w: float }

pub impl VecTransmute: Index<uint, float> {
    #[inline(always)]
    pure fn index(i: uint) -> float unsafe {
        (transmute::<VecTransmute, [float * 4]>(self))[i]
    }
}

pub struct VecBufSlice { x: float, y: float, z: float, w: float }

pub impl VecBufSlice: Index<uint, float> {
    #[inline(always)]
    pure fn index(i: uint) -> float unsafe {
        do vec::raw::buf_as_slice(
            transmute::<*VecBufSlice, *float>(
                to_unsafe_ptr(&self)), 4) |slice| { slice[i] }
    }
}

fn main() {
    let n_tests = 100000;

    let v = Vector { x: 1f, y: 2f, z: 3f, w: 4f };

    let vfield_avg = do test_avg_time_ns(n_tests) {
        // io::println(fmt!("[ %?, %?, %?, %? ]", v.x, v.y, v.z, v.w));
        assert(v.x == 1f);
        assert(v.y == 2f);
        assert(v.z == 3f);
        assert(v.w == 4f);
    };

    io::println(fmt!("Vector.x,y,z,w:      [%d]", vfield_avg as int));

    let vmatch = VecMatch { x: 1f, y: 2f, z: 3f, w: 4f };

    let vindex_avg = do test_avg_time_ns(n_tests) {
        assert(vmatch[0] == 1f);
        assert(vmatch[1] == 2f);
        assert(vmatch[2] == 3f);
        assert(vmatch[3] == 4f);
    };

    let vreint = VecReinterpret { x: 1f, y: 2f, z: 3f, w: 4f };

    let vreint_avg = do test_avg_time_ns(n_tests) {
        assert(vreint[0] == 1f);
        assert(vreint[1] == 2f);
        assert(vreint[2] == 3f);
        assert(vreint[3] == 4f);
    };

    let vtrans = VecTransmute { x: 1f, y: 2f, z: 3f, w: 4f };

    let vtrans_avg = do test_avg_time_ns(n_tests) {
        assert(vtrans[0] == 1f);
        assert(vtrans[1] == 2f);
        assert(vtrans[2] == 3f);
        assert(vtrans[3] == 4f);
    };

    let vbufslice = VecBufSlice { x: 1f, y: 2f, z: 3f, w: 4f };

    let vbufslice_avg = do test_avg_time_ns(n_tests) {
        assert(vbufslice[0] == 1f);
        assert(vbufslice[1] == 2f);
        assert(vbufslice[2] == 3f);
        assert(vbufslice[3] == 4f);
    };

    let min = [vfield_avg, vindex_avg, vreint_avg, vtrans_avg, vbufslice_avg].min();

    io::println(fmt!("Vector.x,y,z,w:    %d = %d", vfield_avg as int, (vfield_avg - min) as int));
    io::println(fmt!("VecMatch[i]:       %d = %d", vindex_avg as int, (vindex_avg - min) as int));
    io::println(fmt!("VecReinterpret[i]: %d = %d", vreint_avg as int, (vreint_avg - min) as int));
    io::println(fmt!("VecTransmute[i]:   %d = %d", vtrans_avg as int, (vtrans_avg - min) as int));
    io::println(fmt!("VecBufSlice[i]:    %d = %d", vbufslice_avg as int, (vbufslice_avg - min) as int));
}

fn test_avg_time_ns(n: uint, f: fn&()) -> u64 {

    let mut total = 0;
    for n.times {
        let start_time = precise_time_ns();

        f();

        total += precise_time_ns() - start_time;
    }

    return total / (n as u64);
}<commit_msg>Add description<commit_after>/**
 * Index operator performance test. For best results, compile with the optimise
 * flag (-O). It seems like VecBufSlice is the fastest version.
 */

extern mod std;
use std::time::precise_time_ns;
use cast::{reinterpret_cast, transmute};
use vec::raw::buf_as_slice;
use ptr::to_unsafe_ptr;

pub struct Vector { x: float, y: float, z: float, w: float }

pub struct VecMatch { x: float, y: float, z: float, w: float }

pub impl VecMatch: Index<uint, float> {
    #[inline(always)]
    pure fn index(i: uint) -> float {
        match i {
            0 => self.x,
            1 => self.y,
            2 => self.z,
            3 => self.w,
            _ => fail(~"Vector index out of bounds.")
        }
    }
}

pub struct VecReinterpret { x: float, y: float, z: float, w: float }

pub impl VecReinterpret: Index<uint, float> {
    #[inline(always)]
    pure fn index(i: uint) -> float unsafe {
        (reinterpret_cast::<VecReinterpret, [float * 4]>(&self))[i]
    }
}

pub struct VecTransmute { x: float, y: float, z: float, w: float }

pub impl VecTransmute: Index<uint, float> {
    #[inline(always)]
    pure fn index(i: uint) -> float unsafe {
        (transmute::<VecTransmute, [float * 4]>(self))[i]
    }
}

pub struct VecBufSlice { x: float, y: float, z: float, w: float }

pub impl VecBufSlice: Index<uint, float> {
    #[inline(always)]
    pure fn index(i: uint) -> float unsafe {
        do vec::raw::buf_as_slice(
            transmute::<*VecBufSlice, *float>(
                to_unsafe_ptr(&self)), 4) |slice| { slice[i] }
    }
}

fn main() {
    let n_tests = 100000;

    let v = Vector { x: 1f, y: 2f, z: 3f, w: 4f };

    let vfield_avg = do test_avg_time_ns(n_tests) {
        // io::println(fmt!("[ %?, %?, %?, %? ]", v.x, v.y, v.z, v.w));
        assert(v.x == 1f);
        assert(v.y == 2f);
        assert(v.z == 3f);
        assert(v.w == 4f);
    };

    io::println(fmt!("Vector.x,y,z,w:      [%d]", vfield_avg as int));

    let vmatch = VecMatch { x: 1f, y: 2f, z: 3f, w: 4f };

    let vindex_avg = do test_avg_time_ns(n_tests) {
        assert(vmatch[0] == 1f);
        assert(vmatch[1] == 2f);
        assert(vmatch[2] == 3f);
        assert(vmatch[3] == 4f);
    };

    let vreint = VecReinterpret { x: 1f, y: 2f, z: 3f, w: 4f };

    let vreint_avg = do test_avg_time_ns(n_tests) {
        assert(vreint[0] == 1f);
        assert(vreint[1] == 2f);
        assert(vreint[2] == 3f);
        assert(vreint[3] == 4f);
    };

    let vtrans = VecTransmute { x: 1f, y: 2f, z: 3f, w: 4f };

    let vtrans_avg = do test_avg_time_ns(n_tests) {
        assert(vtrans[0] == 1f);
        assert(vtrans[1] == 2f);
        assert(vtrans[2] == 3f);
        assert(vtrans[3] == 4f);
    };

    let vbufslice = VecBufSlice { x: 1f, y: 2f, z: 3f, w: 4f };

    let vbufslice_avg = do test_avg_time_ns(n_tests) {
        assert(vbufslice[0] == 1f);
        assert(vbufslice[1] == 2f);
        assert(vbufslice[2] == 3f);
        assert(vbufslice[3] == 4f);
    };

    let min = [vfield_avg, vindex_avg, vreint_avg, vtrans_avg, vbufslice_avg].min();

    io::println(fmt!("Vector.x,y,z,w:    %d = %d", vfield_avg as int, (vfield_avg - min) as int));
    io::println(fmt!("VecMatch[i]:       %d = %d", vindex_avg as int, (vindex_avg - min) as int));
    io::println(fmt!("VecReinterpret[i]: %d = %d", vreint_avg as int, (vreint_avg - min) as int));
    io::println(fmt!("VecTransmute[i]:   %d = %d", vtrans_avg as int, (vtrans_avg - min) as int));
    io::println(fmt!("VecBufSlice[i]:    %d = %d", vbufslice_avg as int, (vbufslice_avg - min) as int));
}

fn test_avg_time_ns(n: uint, f: fn&()) -> u64 {

    let mut total = 0;
    for n.times {
        let start_time = precise_time_ns();

        f();

        total += precise_time_ns() - start_time;
    }

    return total / (n as u64);
}<|endoftext|>


<commit_before>#![allow(non_snake_case)]

extern crate winapi;
extern crate rustup;

use std::ffi::CString;
use std::path::PathBuf;
use ::winapi::{HRESULT, PCSTR, UINT, LPCWSTR, LPWSTR, LPVOID};

pub type MSIHANDLE = u32;

pub const LOGMSG_TRACEONLY: i32 = 0;
pub const LOGMSG_VERBOSE: i32 = 1;
pub const LOGMSG_STANDARD: i32 = 2;

// TODO: share this with self_update.rs
static TOOLS: &'static [&'static str]
    = &["rustc", "rustdoc", "cargo", "rust-lldb", "rust-gdb"];

#[no_mangle]
/// This is run as an `immediate` action early in the install sequence
pub unsafe extern "system" fn RustupSetInstallLocation(hInstall: MSIHANDLE) -> UINT {
    // TODO: error handling (get rid of unwrap)
    let name = CString::new("RustupSetInstallLocation").unwrap();
    let hr = WcaInitialize(hInstall, name.as_ptr());
    //let path = ::rustup::utils::cargo_home().unwrap();
    let path = PathBuf::from(::std::env::var_os("USERPROFILE").unwrap()).join(".rustup-test");
    set_property("RustupInstallLocation", path.to_str().unwrap());
    WcaFinalize(hr)
}

#[no_mangle]
/// This is be run as a `deferred` action after `InstallFiles` on install and upgrade
pub unsafe extern "system" fn RustupInstall(hInstall: MSIHANDLE) -> UINT {
    let name = CString::new("RustupInstall").unwrap();
    let hr = WcaInitialize(hInstall, name.as_ptr());
    // For deferred custom actions, all data must be passed through the `CustomActionData` property
    let custom_action_data = get_property("CustomActionData");
    // TODO: use rustup_utils::cargo_home() or pass through CustomActionData
    let path = PathBuf::from(::std::env::var_os("USERPROFILE").unwrap()).join(".rustup-test");
    let bin_path = path.join("bin");
    let rustup_path = bin_path.join("rustup.exe");
    let exe_installed = rustup_path.exists();
    log(&format!("Hello World from RustupInstall, confirming that rustup.exe has been installed: {}! CustomActionData: {}", exe_installed, custom_action_data));
    for tool in TOOLS {
        let ref tool_path = bin_path.join(&format!("{}.exe", tool));
        ::rustup::utils::hardlink_file(&rustup_path, tool_path);
    }
    // TODO: install default toolchain and report progress to UI
    WcaFinalize(hr)
}

#[no_mangle]
/// This is be run as a `deferred` action after `RemoveFiles` on uninstall (not on upgrade!)
pub unsafe extern "system" fn RustupUninstall(hInstall: MSIHANDLE) -> UINT {
    let name = CString::new("RustupUninstall").unwrap();
    let hr = WcaInitialize(hInstall, name.as_ptr());
    // For deferred custom actions, all data must be passed through the `CustomActionData` property
    let custom_action_data = get_property("CustomActionData");
    // TODO: use rustup_utils::cargo_home() or pass through CustomActionData
    let path = PathBuf::from(::std::env::var_os("USERPROFILE").unwrap()).join(".rustup-test");
    let exe_deleted = !path.join("bin").join("rustup.exe").exists();
    log(&format!("Hello World from RustupUninstall, confirming that rustup.exe has been deleted: {}! CustomActionData: {}", exe_deleted, custom_action_data));
    // TODO: Remove .cargo and .rustup
    ::rustup::utils::remove_dir("rustup-test", &path, &|_| {});
    WcaFinalize(hr)
}

// wrapper for WcaGetProperty (TODO: error handling)
fn get_property(name: &str) -> String {
    let encoded_name = to_wide_chars(name);
    let mut result_ptr = std::ptr::null_mut();
    unsafe { WcaGetProperty(encoded_name.as_ptr(), &mut result_ptr) };
    let result = from_wide_ptr(result_ptr);
    unsafe { StrFree(result_ptr as LPVOID) };
    result
}

// wrapper for WcaSetProperty
fn set_property(name: &str, value: &str) -> HRESULT {
    let encoded_name = to_wide_chars(name);
    let encoded_value = to_wide_chars(value);
    unsafe { WcaSetProperty(encoded_name.as_ptr(), encoded_value.as_ptr()) }
}

fn log(message: &str) {
    let msg = CString::new(message).unwrap();
    unsafe { WcaLog(LOGMSG_STANDARD, msg.as_ptr()) }
}
fn from_wide_ptr(ptr: *const u16) -> String {
    use std::ffi::OsString;
    use std::os::windows::ffi::OsStringExt;
    unsafe {
        assert!(!ptr.is_null());
        let len = (0..std::isize::MAX).position(|i| *ptr.offset(i) == 0).unwrap();
        let slice = std::slice::from_raw_parts(ptr, len);
        OsString::from_wide(slice).to_string_lossy().into_owned()
    }
}

fn to_wide_chars(s: &str) -> Vec<u16> {
    use std::ffi::OsStr;
    use std::os::windows::ffi::OsStrExt;
    OsStr::new(s).encode_wide().chain(Some(0).into_iter()).collect::<Vec<_>>()
}

extern "system" {
    fn WcaInitialize(hInstall: MSIHANDLE, szCustomActionLogName: PCSTR) -> HRESULT;
    fn WcaFinalize(iReturnValue: HRESULT) -> UINT;
    fn WcaGetProperty(wzProperty: LPCWSTR, ppwzData: *mut LPWSTR) -> HRESULT; // see documentation for MsiGetProperty
    fn WcaSetProperty(wzPropertyName: LPCWSTR, wzPropertyValue: LPCWSTR) -> HRESULT;
    fn StrFree(p: LPVOID) -> HRESULT;
}

extern "cdecl" {
    fn WcaLog(llv: i32, fmt: PCSTR);
}<commit_msg>Update outdated list of tools<commit_after>#![allow(non_snake_case)]

extern crate winapi;
extern crate rustup;

use std::ffi::CString;
use std::path::PathBuf;
use ::winapi::{HRESULT, PCSTR, UINT, LPCWSTR, LPWSTR, LPVOID};

pub type MSIHANDLE = u32;

pub const LOGMSG_TRACEONLY: i32 = 0;
pub const LOGMSG_VERBOSE: i32 = 1;
pub const LOGMSG_STANDARD: i32 = 2;

// TODO: share this with self_update.rs
static TOOLS: &'static [&'static str]
    = &["rustc", "rustdoc", "cargo", "rust-lldb", "rust-gdb", "rls"];

#[no_mangle]
/// This is run as an `immediate` action early in the install sequence
pub unsafe extern "system" fn RustupSetInstallLocation(hInstall: MSIHANDLE) -> UINT {
    // TODO: error handling (get rid of unwrap)
    let name = CString::new("RustupSetInstallLocation").unwrap();
    let hr = WcaInitialize(hInstall, name.as_ptr());
    //let path = ::rustup::utils::cargo_home().unwrap();
    let path = PathBuf::from(::std::env::var_os("USERPROFILE").unwrap()).join(".rustup-test");
    set_property("RustupInstallLocation", path.to_str().unwrap());
    WcaFinalize(hr)
}

#[no_mangle]
/// This is be run as a `deferred` action after `InstallFiles` on install and upgrade
pub unsafe extern "system" fn RustupInstall(hInstall: MSIHANDLE) -> UINT {
    let name = CString::new("RustupInstall").unwrap();
    let hr = WcaInitialize(hInstall, name.as_ptr());
    // For deferred custom actions, all data must be passed through the `CustomActionData` property
    let custom_action_data = get_property("CustomActionData");
    // TODO: use rustup_utils::cargo_home() or pass through CustomActionData
    let path = PathBuf::from(::std::env::var_os("USERPROFILE").unwrap()).join(".rustup-test");
    let bin_path = path.join("bin");
    let rustup_path = bin_path.join("rustup.exe");
    let exe_installed = rustup_path.exists();
    log(&format!("Hello World from RustupInstall, confirming that rustup.exe has been installed: {}! CustomActionData: {}", exe_installed, custom_action_data));
    for tool in TOOLS {
        let ref tool_path = bin_path.join(&format!("{}.exe", tool));
        ::rustup::utils::hardlink_file(&rustup_path, tool_path);
    }
    // TODO: install default toolchain and report progress to UI
    WcaFinalize(hr)
}

#[no_mangle]
/// This is be run as a `deferred` action after `RemoveFiles` on uninstall (not on upgrade!)
pub unsafe extern "system" fn RustupUninstall(hInstall: MSIHANDLE) -> UINT {
    let name = CString::new("RustupUninstall").unwrap();
    let hr = WcaInitialize(hInstall, name.as_ptr());
    // For deferred custom actions, all data must be passed through the `CustomActionData` property
    let custom_action_data = get_property("CustomActionData");
    // TODO: use rustup_utils::cargo_home() or pass through CustomActionData
    let path = PathBuf::from(::std::env::var_os("USERPROFILE").unwrap()).join(".rustup-test");
    let exe_deleted = !path.join("bin").join("rustup.exe").exists();
    log(&format!("Hello World from RustupUninstall, confirming that rustup.exe has been deleted: {}! CustomActionData: {}", exe_deleted, custom_action_data));
    // TODO: Remove .cargo and .rustup
    ::rustup::utils::remove_dir("rustup-test", &path, &|_| {});
    WcaFinalize(hr)
}

// wrapper for WcaGetProperty (TODO: error handling)
fn get_property(name: &str) -> String {
    let encoded_name = to_wide_chars(name);
    let mut result_ptr = std::ptr::null_mut();
    unsafe { WcaGetProperty(encoded_name.as_ptr(), &mut result_ptr) };
    let result = from_wide_ptr(result_ptr);
    unsafe { StrFree(result_ptr as LPVOID) };
    result
}

// wrapper for WcaSetProperty
fn set_property(name: &str, value: &str) -> HRESULT {
    let encoded_name = to_wide_chars(name);
    let encoded_value = to_wide_chars(value);
    unsafe { WcaSetProperty(encoded_name.as_ptr(), encoded_value.as_ptr()) }
}

fn log(message: &str) {
    let msg = CString::new(message).unwrap();
    unsafe { WcaLog(LOGMSG_STANDARD, msg.as_ptr()) }
}
fn from_wide_ptr(ptr: *const u16) -> String {
    use std::ffi::OsString;
    use std::os::windows::ffi::OsStringExt;
    unsafe {
        assert!(!ptr.is_null());
        let len = (0..std::isize::MAX).position(|i| *ptr.offset(i) == 0).unwrap();
        let slice = std::slice::from_raw_parts(ptr, len);
        OsString::from_wide(slice).to_string_lossy().into_owned()
    }
}

fn to_wide_chars(s: &str) -> Vec<u16> {
    use std::ffi::OsStr;
    use std::os::windows::ffi::OsStrExt;
    OsStr::new(s).encode_wide().chain(Some(0).into_iter()).collect::<Vec<_>>()
}

extern "system" {
    fn WcaInitialize(hInstall: MSIHANDLE, szCustomActionLogName: PCSTR) -> HRESULT;
    fn WcaFinalize(iReturnValue: HRESULT) -> UINT;
    fn WcaGetProperty(wzProperty: LPCWSTR, ppwzData: *mut LPWSTR) -> HRESULT; // see documentation for MsiGetProperty
    fn WcaSetProperty(wzPropertyName: LPCWSTR, wzPropertyValue: LPCWSTR) -> HRESULT;
    fn StrFree(p: LPVOID) -> HRESULT;
}

extern "cdecl" {
    fn WcaLog(llv: i32, fmt: PCSTR);
}<|endoftext|>