<commit_before>#[derive(Debug)]
enum Direction {
    Diagonal,
    Up,
    Left,
    Undef,
}

impl Clone for Direction {
    fn clone(&self) -> Direction {
        match *self {
            Diagonal => Direction::Diagonal,
            Up => Direction::Up,
            Left => Direction::Left,
            Undef => Direction::Undef,
        }
    }
}

#[derive(Debug, Clone)]
struct Score<T> {
    val: T,
    dir: Direction,
}

impl Score {
    fn new() -> Score {
        Score { val: 0, dir: Undef }
    }
}

// mod tps_tree;

// use tps_tree::{TpsTree, TpsJoint};
use Direction::*;
use std::cmp::{max, min};

fn get_hsp(w1: &String, w2: &String, match_scr: i32, mism_scr: i32, gap_scr: i32) -> i32 {
    let mut hsp = 0;

    let mut chars2 = w2.chars();
    for c1 in w1.chars() {
        let c2 = chars2.next().unwrap();
        hsp += if c1 == c2 {
            match_scr
        } else if c1 == '_' || c2 == '_' {
            gap_scr
        } else {
            mism_scr
        }
    }

    hsp
}

fn print_matrix(seq_vec1: &Vec<String>, seq_vec2: &Vec<String>, mtrx: &Vec<Vec<Score>>) {
    let len_2 = seq_vec2.len();
    for sec in seq_vec1.iter() {
        for _ in 0..len_2 {
            print!("\t");
        }
        for chr in sec.chars() {
            print!("{}\t", chr);
        }
        println!("");
    }

    for i in 0..seq_vec2[0].len() {
        for sec in seq_vec2.iter() {
            print!("{}\t", sec.chars().nth(i).unwrap());
        }
        for elem in mtrx[i].iter() {
            print!("{}\t", elem.val);
        }
        println!("");
    }
}

fn align_seqs_seq(mut seq_vec: Vec<String>, mut seq: String) {
    /********* Initialize sequences **********/
    seq = String::from("_") + seq.as_ref();

    for sec in seq_vec.iter_mut() {
        *sec = String::from("_") + sec.as_ref();
    }
    /********* score function **********/
    fn score(a: char, b: char) -> i32 {
        if a == b {
            0
        } else if a == '_' || b == '_' {
            2
        } else {
            3
        }
    }
    /********* Initialize matrix **********/
    let len_1 = seq_vec[0].len();
    let len_2 = seq.len();

    let mut matrix: Vec<Vec<Score>> = vec![vec![Score::new(); len_1]; len_2];

    let mut val = 0;
    for j in 1..len_1 {
        for sec in seq_vec.iter() {
            val += score(seq.chars().nth(0).unwrap(), sec.chars().nth(j).unwrap());
        }
        matrix[0][j].val = val;
        matrix[0][j].dir = Left;
    }

    val = 0;
    for i in 1..len_2 {
        for sec in seq_vec.iter() {
            val += score(seq.chars().nth(i).unwrap(), sec.chars().nth(0).unwrap());
        }
        matrix[i][0].val = val;
        matrix[i][0].dir = Up;
    }
    /**************************************/

    let mut diag_val;
    let mut left_val;
    let mut uppr_val;

    for i in 1..len_2 {
        for j in 1..len_1 {
            diag_val = {
                let mut sum = 0;
                for sec in seq_vec.iter() {
                    sum += score(seq.chars().nth(i).unwrap(), sec.chars().nth(j).unwrap());
                }
                sum += matrix[i - 1][j - 1].val;
                sum
            };

            left_val = {
                let mut sum = 0;
                for _ in seq_vec.iter() {
                    sum += score(seq.chars().nth(i).unwrap(), '_');
                }
                sum += matrix[i][j - 1].val;
                sum
            };

            uppr_val = {
                let mut sum = 0;
                for sec in seq_vec.iter() {
                    sum += score('_', sec.chars().nth(j).unwrap());
                }
                sum += matrix[i - 1][j].val;
                sum
            };

            matrix[i][j].val = min(diag_val, min(left_val, uppr_val));
            matrix[i][j].dir = if matrix[i][j].val == diag_val {
                Diagonal
            } else if matrix[i][j].val == left_val {
                Left
            } else if matrix[i][j].val == uppr_val {
                Up
            } else {
                Undef
            };
        }
    }

    print_matrix(&seq_vec, &seq, &matrix);

    /********* Restore sequences **********/
    seq.remove(0);

    for sec in seq_vec.iter_mut() {
        sec.remove(0);
    }
    /**************************************/

    let mut i = len_2 - 1;
    let mut j = len_1 - 1;
    let mut dir: Direction;
    while i != 0 && j != 0 {
        dir = matrix[i][j].dir.clone();

        match dir {
            Diagonal => {
                i -= 1;
                j -= 1;
            }
            Left => {
                let term = seq.split_off(j);
                seq = format!("{}_{}", seq, term);
                j -= 1;
            }
            Up => {
                for mut sec in seq_vec.iter_mut() {
                    let term = sec.split_off(i);
                    *sec = format!("{}_{}", sec, term);
                }
                i -= 1;
            }
            Undef => {
                break;
            }
        }
    }

    println!("\n> {}", seq);
    for sec in seq_vec.iter() {
        println!("> {}", sec)
    }
}

fn align_alignments(mut seq_vec1: Vec<String>, mut seq_vec2: Vec<String>) {
    /********* Initialize sequences **********/
    for sec in seq_vec1.iter_mut() {
        *sec = String::from("_") + sec.as_ref();
    }
    for sec in seq_vec2.iter_mut() {
        *sec = String::from("_") + sec.as_ref();
    }
    /********* score function **********/
    fn score(a: char, b: char) -> f32 {
        if a == '_' || b == '_' {
            0.0
        } else if a == b {
            2.0
        } else {
            1.0
        }
    }
    /********* Initialize matrix **********/
    let len_1 = seq_vec1[0].len();
    let len_2 = seq_vec2[0].len();
    let n_seqs = (seq_vec1.len() + seq_vec2.len()) as f32;

    let mut matrix: Vec<Vec<Score>> = vec![vec![Score::new(); len_1]; len_2];

    for j in 1..len_1 {
        matrix[0][j].dir = Left;
    }
    for i in 1..len_2 {
        matrix[i][0].dir = Up;
    }
    /**************************************/
    let mut diag_val;
    let mut left_val;
    let mut uppr_val;

    for i in 1..len_2 {
        for j in 1..len_1 {
            // println!("({}, {})", j, i);
            diag_val = {
                let mut sum = 0.0;
                for sec1 in seq_vec1.iter() {
                    for sec2 in seq_vec2.iter() {
                        sum += score(sec1.chars().nth(j).unwrap(), sec2.chars().nth(i).unwrap());
                        // println!("{}:{}", sec1.chars().nth(j).unwrap(), sec2.chars().nth(i).unwrap());
                    }
                }
                sum /= n_seqs;
                sum += matrix[i - 1][j - 1].val;
                // println!("D: {}\n", sum);
                sum
            };

            left_val = {
                let mut sum = 0.0;
                for sec1 in seq_vec1.iter() {
                    for sec2 in seq_vec2.iter() {
                        sum += score(sec1.chars().nth(j-1).unwrap(), sec2.chars().nth(i).unwrap());
                        // println!("{}:{}", sec1.chars().nth(j-1).unwrap(), sec2.chars().nth(i).unwrap());
                    }
                }
                sum /= n_seqs;
                // sum += matrix[i][j - 1].val;
                // println!("L: {}\n", sum);
                sum
            };

            uppr_val = {
                let mut sum = 0.0;
                for sec1 in seq_vec1.iter() {
                    for sec2 in seq_vec2.iter() {
                        sum += score(sec1.chars().nth(j).unwrap(), sec2.chars().nth(i-1).unwrap());
                        // println!("{}:{}", sec1.chars().nth(j).unwrap(), sec2.chars().nth(i-1).unwrap());
                    }
                }
                sum /= n_seqs;
                // sum += matrix[i - 1][j].val;
                // println!("U: {}\n", sum);
                sum
            };

            if diag_val >= left_val {
                if diag_val >= uppr_val {
                    matrix[i][j].val = diag_val;
                    matrix[i][j].dir = Diagonal;
                } else {
                    matrix[i][j].val = uppr_val;
                    matrix[i][j].dir = Up;
                }
            } else {
                if left_val >= uppr_val {
                    matrix[i][j].val = left_val;
                    matrix[i][j].dir = Left;
                } else {
                    matrix[i][j].val = uppr_val;
                    matrix[i][j].dir = Up;
                }
            }
        }
    }

    print_matrix(&seq_vec1, &seq_vec2, &matrix);

    /********* Restore sequences **********/
    for sec in seq_vec1.iter_mut() {
        sec.remove(0);
    }
    for sec in seq_vec2.iter_mut() {
        sec.remove(0);
    }
    /**************************************/

    let mut i = len_2 - 1;
    let mut j = len_1 - 1;
    let mut dir: Direction;
    while i != 0 || j != 0 {
        dir = matrix[i][j].dir.clone();
        println!("{:?}", dir);

        match dir {
            Diagonal => {
                i -= 1;
                j -= 1;
            }
            Left => {
                for mut sec in seq_vec2.iter_mut() {
                    let term = sec.split_off(i);
                    *sec = format!("{}_{}", sec, term);
                }
                j -= 1;
            }
            Up => {
                for mut sec in seq_vec1.iter_mut() {
                    let term = sec.split_off(j);
                    *sec = format!("{}_{}", sec, term);
                }
                i -= 1;
            }
            Undef => {
                break;
            }
        }
    }

    for sec in seq_vec1.iter() {
        println!("> {}", sec)
    }
    for sec in seq_vec2.iter() {
        println!("> {}", sec)
    }
}

fn pair_sum(sec_1: &str, sec_2: &str, match_scr: i32, mism_scr: i32, gap_scr: i32) -> f32 {
    let len_1 = sec_1.len();
    let len_2 = sec_2.len();
    let diff: i32 = len_1 as i32 - len_2 as i32;

    let sec_1s;
    let sec_2s;

    if diff > 0 {
        sec_1s = String::from(sec_1);
        sec_2s = format!("{}{}", sec_2, "_".repeat(diff as usize));
    } else {
        sec_1s = format!("{}{}", sec_1, "_".repeat(diff.abs() as usize));
        sec_2s = String::from(sec_2);
    }
    // println!("A: {}\nB: {}\n> {}\n", sec_1s, sec_2s, get_hsp(&sec_1s, &sec_2s, match_scr, mism_scr, gap_scr) as f32 / sec_1s.len() as f32);
    get_hsp(&sec_1s, &sec_2s, match_scr, mism_scr, gap_scr) as f32 / sec_1s.len() as f32
}

pub fn tps_alignment(seqs: Vec<&str>) {
    let len = seqs.len();
    let mut my_vec = Vec::new();

    for i in 0..len {
        let mut j = i + 1;
        while j < len {
            my_vec.push((pair_sum(seqs[i], seqs[j], 2, 1, 1), i, j));
            j += 1;
        }
    }

    my_vec.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap());
    // println!("{:?}", my_vec);

    align_alignments(vec![String::from("AAAC"), String::from("_GAC")],
                     vec![String::from("AGC"), String::from("ACC")]);

    align_seqs_seq(vec![String::from("ACTCAT"), String::from("AGTCAT")],
                   String::from("ACGTCCT"));

    // align_alignments(vec![String::from("TGTAAC")],
    //                  vec![String::from("TGTAC")]);

    // let tree = TpsTree::new();
    // let joint1 = TpsJoint::from_strings(seqs[my_vec[0].1], seqs[my_vec[0].2]);
    // println!("{:?}", joint1);
    // let joint2 = TpsJoint::from_strings(seqs[my_vec[2].1], seqs[my_vec[2].2]);
    // println!("{:?}", joint2);
}<commit_msg>print<T> working!<commit_after>#[derive(Debug)]
enum Direction {
    Diagonal,
    Up,
    Left,
    Undef,
}

impl Clone for Direction {
    fn clone(&self) -> Direction {
        match *self {
            Diagonal => Direction::Diagonal,
            Up => Direction::Up,
            Left => Direction::Left,
            Undef => Direction::Undef,
        }
    }
}

#[derive(Debug, Clone)]
struct Score<T> {
    val: T,
    dir: Direction,
}

impl Score<i32> {
    fn new() -> Score<i32> {
        Score { val: 0, dir: Undef }
    }
}

impl Score<f32> {
    fn new() -> Score<f32> {
        Score { val: 0.0, dir: Undef }
    }
}

// mod tps_tree;

// use tps_tree::{TpsTree, TpsJoint};
use Direction::*;
use std::cmp::{max, min};

fn get_hsp(w1: &String, w2: &String, match_scr: i32, mism_scr: i32, gap_scr: i32) -> i32 {
    let mut hsp = 0;

    let mut chars2 = w2.chars();
    for c1 in w1.chars() {
        let c2 = chars2.next().unwrap();
        hsp += if c1 == c2 {
            match_scr
        } else if c1 == '_' || c2 == '_' {
            gap_scr
        } else {
            mism_scr
        }
    }

    hsp
}

fn print_matrix<T>(seq_vec1: &Vec<String>, seq_vec2: &Vec<String>, mtrx: &Vec<Vec<Score<T>>>)
    where T: std::fmt::Display {
    let len_2 = seq_vec2.len();
    for sec in seq_vec1.iter() {
        for _ in 0..len_2 {
            print!("\t");
        }
        for chr in sec.chars() {
            print!("{}\t", chr);
        }
        println!("");
    }

    for i in 0..seq_vec2[0].len() {
        for sec in seq_vec2.iter() {
            print!("{}\t", sec.chars().nth(i).unwrap());
        }
        for elem in mtrx[i].iter() {
            print!("{}\t", elem.val);
        }
        println!("");
    }
}

fn align_seqs_seq(mut seq_vec: Vec<String>, mut seq: String) {
    /********* Initialize sequences **********/
    seq = String::from("_") + seq.as_ref();

    for sec in seq_vec.iter_mut() {
        *sec = String::from("_") + sec.as_ref();
    }
    /********* score function **********/
    fn score(a: char, b: char) -> i32 {
        if a == b {
            0
        } else if a == '_' || b == '_' {
            2
        } else {
            3
        }
    }
    /********* Initialize matrix **********/
    let len_1 = seq_vec[0].len();
    let len_2 = seq.len();

    let mut matrix: Vec<Vec<Score<i32>>> = vec![vec![Score::<i32>::new(); len_1]; len_2];

    let mut val = 0;
    for j in 1..len_1 {
        for sec in seq_vec.iter() {
            val += score(seq.chars().nth(0).unwrap(), sec.chars().nth(j).unwrap());
        }
        matrix[0][j].val = val;
        matrix[0][j].dir = Left;
    }

    val = 0;
    for i in 1..len_2 {
        for sec in seq_vec.iter() {
            val += score(seq.chars().nth(i).unwrap(), sec.chars().nth(0).unwrap());
        }
        matrix[i][0].val = val;
        matrix[i][0].dir = Up;
    }
    /**************************************/

    let mut diag_val;
    let mut left_val;
    let mut uppr_val;

    for i in 1..len_2 {
        for j in 1..len_1 {
            diag_val = {
                let mut sum = 0;
                for sec in seq_vec.iter() {
                    sum += score(seq.chars().nth(i).unwrap(), sec.chars().nth(j).unwrap());
                }
                sum += matrix[i - 1][j - 1].val;
                sum
            };

            left_val = {
                let mut sum = 0;
                for _ in seq_vec.iter() {
                    sum += score(seq.chars().nth(i).unwrap(), '_');
                }
                sum += matrix[i][j - 1].val;
                sum
            };

            uppr_val = {
                let mut sum = 0;
                for sec in seq_vec.iter() {
                    sum += score('_', sec.chars().nth(j).unwrap());
                }
                sum += matrix[i - 1][j].val;
                sum
            };

            matrix[i][j].val = min(diag_val, min(left_val, uppr_val));
            matrix[i][j].dir = if matrix[i][j].val == diag_val {
                Diagonal
            } else if matrix[i][j].val == left_val {
                Left
            } else if matrix[i][j].val == uppr_val {
                Up
            } else {
                Undef
            };
        }
    }

    print_matrix(&seq_vec, &vec![seq.clone()], &matrix);

    /********* Restore sequences **********/
    seq.remove(0);

    for sec in seq_vec.iter_mut() {
        sec.remove(0);
    }
    /**************************************/

    let mut i = len_2 - 1;
    let mut j = len_1 - 1;
    let mut dir: Direction;
    while i != 0 && j != 0 {
        dir = matrix[i][j].dir.clone();

        match dir {
            Diagonal => {
                i -= 1;
                j -= 1;
            }
            Left => {
                let term = seq.split_off(j);
                seq = format!("{}_{}", seq, term);
                j -= 1;
            }
            Up => {
                for mut sec in seq_vec.iter_mut() {
                    let term = sec.split_off(i);
                    *sec = format!("{}_{}", sec, term);
                }
                i -= 1;
            }
            Undef => {
                break;
            }
        }
    }

    println!("\n> {}", seq);
    for sec in seq_vec.iter() {
        println!("> {}", sec)
    }
}

fn align_alignments(mut seq_vec1: Vec<String>, mut seq_vec2: Vec<String>) {
    /********* Initialize sequences **********/
    for sec in seq_vec1.iter_mut() {
        *sec = String::from("_") + sec.as_ref();
    }
    for sec in seq_vec2.iter_mut() {
        *sec = String::from("_") + sec.as_ref();
    }
    /********* score function **********/
    fn score(a: char, b: char) -> f32 {
        if a == '_' || b == '_' {
            0.0
        } else if a == b {
            2.0
        } else {
            1.0
        }
    }
    /********* Initialize matrix **********/
    let len_1 = seq_vec1[0].len();
    let len_2 = seq_vec2[0].len();
    let n_seqs = (seq_vec1.len() + seq_vec2.len()) as f32;

    let mut matrix: Vec<Vec<Score<f32>>> = vec![vec![Score::<f32>::new(); len_1]; len_2];

    for j in 1..len_1 {
        matrix[0][j].dir = Left;
    }
    for i in 1..len_2 {
        matrix[i][0].dir = Up;
    }
    /**************************************/
    let mut diag_val;
    let mut left_val;
    let mut uppr_val;

    for i in 1..len_2 {
        for j in 1..len_1 {
            // println!("({}, {})", j, i);
            diag_val = {
                let mut sum = 0.0;
                for sec1 in seq_vec1.iter() {
                    for sec2 in seq_vec2.iter() {
                        sum += score(sec1.chars().nth(j).unwrap(), sec2.chars().nth(i).unwrap());
                        // println!("{}:{}", sec1.chars().nth(j).unwrap(), sec2.chars().nth(i).unwrap());
                    }
                }
                sum /= n_seqs;
                sum += matrix[i - 1][j - 1].val;
                // println!("D: {}\n", sum);
                sum
            };

            left_val = {
                let mut sum = 0.0;
                for sec1 in seq_vec1.iter() {
                    for sec2 in seq_vec2.iter() {
                        sum += score(sec1.chars().nth(j-1).unwrap(), sec2.chars().nth(i).unwrap());
                        // println!("{}:{}", sec1.chars().nth(j-1).unwrap(), sec2.chars().nth(i).unwrap());
                    }
                }
                sum /= n_seqs;
                // sum += matrix[i][j - 1].val;
                // println!("L: {}\n", sum);
                sum
            };

            uppr_val = {
                let mut sum = 0.0;
                for sec1 in seq_vec1.iter() {
                    for sec2 in seq_vec2.iter() {
                        sum += score(sec1.chars().nth(j).unwrap(), sec2.chars().nth(i-1).unwrap());
                        // println!("{}:{}", sec1.chars().nth(j).unwrap(), sec2.chars().nth(i-1).unwrap());
                    }
                }
                sum /= n_seqs;
                // sum += matrix[i - 1][j].val;
                // println!("U: {}\n", sum);
                sum
            };

            if diag_val >= left_val {
                if diag_val >= uppr_val {
                    matrix[i][j].val = diag_val;
                    matrix[i][j].dir = Diagonal;
                } else {
                    matrix[i][j].val = uppr_val;
                    matrix[i][j].dir = Up;
                }
            } else {
                if left_val >= uppr_val {
                    matrix[i][j].val = left_val;
                    matrix[i][j].dir = Left;
                } else {
                    matrix[i][j].val = uppr_val;
                    matrix[i][j].dir = Up;
                }
            }
        }
    }

    print_matrix(&seq_vec1, &seq_vec2, &matrix);

    /********* Restore sequences **********/
    for sec in seq_vec1.iter_mut() {
        sec.remove(0);
    }
    for sec in seq_vec2.iter_mut() {
        sec.remove(0);
    }
    /**************************************/

    let mut i = len_2 - 1;
    let mut j = len_1 - 1;
    let mut dir: Direction;
    while i != 0 || j != 0 {
        dir = matrix[i][j].dir.clone();
        match dir {
            Diagonal => {
                i -= 1;
                j -= 1;
            }
            Left => {
                for mut sec in seq_vec2.iter_mut() {
                    let term = sec.split_off(i);
                    *sec = format!("{}_{}", sec, term);
                }
                j -= 1;
            }
            Up => {
                for mut sec in seq_vec1.iter_mut() {
                    let term = sec.split_off(j);
                    *sec = format!("{}_{}", sec, term);
                }
                i -= 1;
            }
            Undef => {
                break;
            }
        }
    }

    for sec in seq_vec1.iter() {
        println!("> {}", sec)
    }
    for sec in seq_vec2.iter() {
        println!("> {}", sec)
    }
}

fn pair_sum(sec_1: &str, sec_2: &str, match_scr: i32, mism_scr: i32, gap_scr: i32) -> f32 {
    let len_1 = sec_1.len();
    let len_2 = sec_2.len();
    let diff: i32 = len_1 as i32 - len_2 as i32;

    let sec_1s;
    let sec_2s;

    if diff > 0 {
        sec_1s = String::from(sec_1);
        sec_2s = format!("{}{}", sec_2, "_".repeat(diff as usize));
    } else {
        sec_1s = format!("{}{}", sec_1, "_".repeat(diff.abs() as usize));
        sec_2s = String::from(sec_2);
    }
    // println!("A: {}\nB: {}\n> {}\n", sec_1s, sec_2s, get_hsp(&sec_1s, &sec_2s, match_scr, mism_scr, gap_scr) as f32 / sec_1s.len() as f32);
    get_hsp(&sec_1s, &sec_2s, match_scr, mism_scr, gap_scr) as f32 / sec_1s.len() as f32
}

pub fn tps_alignment(seqs: Vec<&str>) {
    let len = seqs.len();
    let mut my_vec = Vec::new();

    for i in 0..len {
        let mut j = i + 1;
        while j < len {
            my_vec.push((pair_sum(seqs[i], seqs[j], 2, 1, 1), i, j));
            j += 1;
        }
    }

    my_vec.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap());
    // println!("{:?}", my_vec);

    align_alignments(vec![String::from("AAAC"), String::from("_GAC")],
                     vec![String::from("AGC"), String::from("ACC")]);

    align_seqs_seq(vec![String::from("ACTCAT"), String::from("AGTCAT")],
                   String::from("ACGTCCT"));

    // align_alignments(vec![String::from("TGTAAC")],
    //                  vec![String::from("TGTAC")]);

    // let tree = TpsTree::new();
    // let joint1 = TpsJoint::from_strings(seqs[my_vec[0].1], seqs[my_vec[0].2]);
    // println!("{:?}", joint1);
    // let joint2 = TpsJoint::from_strings(seqs[my_vec[2].1], seqs[my_vec[2].2]);
    // println!("{:?}", joint2);
}<|endoftext|>


<commit_before>use anyhow::Result;
use clap::Parser;
use regex::Regex;
use std::collections::HashMap;
use std::fs::File;
use std::io::{self, BufRead};
use std::path::PathBuf;
use std::time::Instant;

use polars::frame::DataFrame;
use polars::prelude::*;

#[derive(Parser, Debug)]
#[clap(author, version, about, long_about = None)]
struct Args {
    /// List of input files. Ideally with headers
    #[clap(short, long, multiple_values(true))]
    input: Vec<PathBuf>,

    /// An output file
    #[clap(short, long)]
    output: String,

    /// The column to select. If files have no headers, use "column_N"
    #[clap(short, long, multiple_values(true))]
    select: Vec<String>,

    /// The size of chunks to process at once
    #[clap(short, long, default_value_t = 100)]
    chunksize: usize,
}

fn load_data(path: &PathBuf, delim: u8) -> Result<DataFrame> {
    let mut df: DataFrame = CsvReader::from_path(path)?
        .has_header(true)
        .with_delimiter(delim)
        .with_null_values(Some(NullValues::AllColumns("NA".to_string())))
        .infer_schema(None)
        .finish()
        .unwrap();

    // hstack the experiment name (derived from the filename) into the DataFrame
    let iter_exp =
        vec![path.file_name().unwrap().to_str().unwrap().split('-').collect::<Vec<&str>>()[0..=2]
            .join("-")]
        .into_iter();
    let exp_col: Series = Series::new(
        "experiment",
        iter_exp.flat_map(|n| std::iter::repeat(n).take(df.height())).collect::<Vec<String>>(),
    );
    // exp_col.rename("experiment");
    df.hstack_mut(&[exp_col]).unwrap();

    Ok(df)
}

fn baseline_get_median_gt_zero(str_val: &Series) -> Series {
    let lists = str_val.utf8().unwrap().into_iter().map(|x| {
        x.unwrap().split(',').into_iter().map(|y| y.parse::<f64>().unwrap()).collect::<Vec<f64>>()
    });

    let medians: Vec<bool> =
        lists.into_iter().map(|x| Series::from_iter(x).median().unwrap() > 0.0).collect();
    Series::from_iter(medians)
}

fn filter_input(input: &mut DataFrame, baseline: bool) -> DataFrame {
    if baseline {
        filter_baseline(input)
    } else {
        filter_differential(input)
    }
}

fn filter_baseline(input: &mut DataFrame) -> DataFrame {
    // This is a baseline experiment.
    // Find columns starting with lower case g, then apply the function to convert to
    // medain and select greater than zero
    let mut meas = Vec::<Expr>::new();

    for column in input.get_column_names_owned() {
        if column.starts_with('g') {
            input.apply(&column, baseline_get_median_gt_zero).unwrap();
            meas.push(col(&column));
        }
    }
    // Selection should now have all the gN column names in it
    input
        .clone()
        .lazy()
        .filter(any_exprs(&meas[0..meas.len()/2]).or(any_exprs(&meas[meas.len()/2 .. meas.len()])) ) // If we try to do the whole thing at once, we get a stack overflow
        .collect()
        .unwrap()
}

fn filter_differential(input: &DataFrame) -> DataFrame {
    // find the p value and log fold columns
    let pv_regex = Regex::new(r".*p-value.*").unwrap();
    let log_fold_regex = Regex::new(r".*log2.*").unwrap();
    let mut exprs: Vec<Expr> = Vec::new();
    let mut pv: Expr;
    let mut lf: Expr;
    for column in input.get_column_names_owned() {
        if pv_regex.is_match(&column) {
            pv = col(&column).is_not_null();
            exprs.push(pv);
        } else if log_fold_regex.is_match(&column) {
            lf = col(&column).neq(lit(0));
            exprs.push(lf);
        }
    }

    input.clone().lazy().filter(all_exprs(exprs)).collect().unwrap()
}

fn load_chunk(
    paths: &mut Vec<PathBuf>,
    select: &mut Vec<String>,
) -> Result<DataFrame, anyhow::Error> {
    /*
    Work with the chunks of input to reduce them into a single dataframe of the genes we want from
    that particular set of experiments. Should return a single dataframe
    */
    let mut output: Option<DataFrame> = None;
    while let Some(infile) = paths.pop() {
        // Load the input Csv
        let mut input: DataFrame = if infile.extension().unwrap() == "tsv" {
            load_data(&infile, b'\t')
        } else {
            load_data(&infile, b',')
        }
        .unwrap_or_else(|error| match error.downcast_ref::<PolarsError>() {
            Some(PolarsError::Io(_string)) => panic!("Input file does not exist! {:?}", infile),
            _ => panic!("An error occurred! {:?}", error),
        });

        // Rename columns to remove . in the names. Now also remove spaces
        let mut new_cols = Vec::new();
        let mut c: u8 = 0;
        for nm in input.get_column_names().iter() {
            let mut new_name = nm.replace('.', "").replace(' ', "");
            if new_cols.contains(&new_name) {
                c += 1;
                new_name += &c.to_string(); // Hopefully avoid duplicate names
            }
            new_cols.push(new_name);
        }

        if new_cols != input.get_column_names() {
            input.set_column_names(&new_cols).unwrap_or_else(|error| {
                panic!("Failed to set column names for some reason {:?}", error);
            });
        }

        // Need to detect which type of experiment it is - differential or baseline
        let type_re = Regex::new(r"tpms").unwrap();
        let filtered = filter_input(&mut input, type_re.is_match(infile.to_str().unwrap()));
        // Drop everything from the input except the genes
        let genes: DataFrame = filtered.select(select.iter()).unwrap_or_else(|error| match error {
            PolarsError::NotFound(_string) => {
                panic!(
                    "{:?} was not found in the header, does the file have a header?\n{:?}",
                    select,
                    input.get_column_names()
                );
            },
            _ => panic!("Error selecting column from input: {:?}", error),
        });

        output = match output {
            None => {
                let output_df = genes.clone();
                Some(output_df)
            },
            Some(mut output_df) => {
                // genes =/= output, so we are handling a new file
                output_df.extend(&genes).unwrap_or_else(|error| {
                    panic!("Failed to extend the output dataframe. {:?}", error)
                });
                output_df =
                    output_df.unique(None, UniqueKeepStrategy::First).unwrap_or_else(|error| {
                        panic!("Failed to parse selected column for uniqueness. {:?}", error)
                    });
                Some(output_df)
            },
        }
    }
    Ok(output.unwrap())
}

fn get_taxid(paths: &mut Vec<PathBuf>, lookup_table: &mut HashMap<String, String>) {
    let tax_regex = Regex::new(r".*Taxon_([0-9]{4,})").unwrap();

    while let Some(path) = paths.pop() {
        let exp_name =
            path.file_name().unwrap().to_str().unwrap().split('-').collect::<Vec<&str>>()[0..=2]
                .join("-")
                .replace(".condensed", "");

        let file = File::open(path).unwrap();
        let lines = io::BufReader::new(file).lines();
        for line in lines {
            match line {
                Ok(line) => {
                    let captures = tax_regex.captures(line.as_str());
                    match captures {
                        Some(caps) => {
                            let taxid = caps.get(1).map_or("", |m| m.as_str());
                            lookup_table.insert(exp_name, taxid.to_string());
                            break;
                        },
                        None => {},
                    };
                },
                Err(_idc) => {},
            }
        }
    }
}

fn main() {
    let timer = Instant::now();

    let mut args = Args::parse();

    // Always select the experiment
    args.select.push("experiment".to_string());

    // separate expression data from taxa data
    let mut taxa_files: Vec<PathBuf> = Vec::new();
    let mut expr_files: Vec<PathBuf> = Vec::new();

    let re = Regex::new(r"sdrf.tsv").unwrap();

    for infile in args.input {
        if re.is_match(infile.to_str().unwrap()) {
            // File is experiment setup data, parse for taxa
            taxa_files.push(infile);
        } else {
            // Should only be expression data, parse for genes
            expr_files.push(infile);
        }
    }

    println!("{} taxa files to parse", taxa_files.len());
    println!("{} expression files to parse", expr_files.len());

    let n_files = expr_files.len();

    let mut dataframes: Vec<DataFrame> = Vec::new();

    let chunked_input: Vec<Vec<PathBuf>> =
        expr_files.chunks(args.chunksize).map(|x| x.to_vec()).collect();
    // Read everything into a big vector
    let mut n_chunks = 0;
    for mut files_chunk in chunked_input {
        dataframes.push(
            load_chunk(&mut files_chunk, &mut args.select)
                .unwrap_or_else(|_error| panic!("Something wrong in one of the reads, aborting")),
        );
        n_chunks += 1;
        println!("Done {} files", n_chunks * args.chunksize);
    }

    let mut output: DataFrame = dataframes[0].clone();

    for item in dataframes.iter() {
        output.extend(item).unwrap_or_else(|error| {
            panic!("Unable to extend the DataFrame! Out of memory? {:?}", error)
        });
        output.unique(None, UniqueKeepStrategy::First).unwrap();
    }

    // Now parse the experiment files to get taxa
    let chunked_taxa: Vec<Vec<PathBuf>> =
        taxa_files.chunks(args.chunksize).map(|x| x.to_vec()).collect();

    let mut lookup_table: HashMap<String, String> = HashMap::new();

    for mut tax_chunk in chunked_taxa {
        get_taxid(&mut tax_chunk, &mut lookup_table);
    }

    // Now we have the experiment - taxid lookup_table, we just need to add the column to the output df

    let expt_col = output.select(["experiment"]).unwrap();
    let mut tax_ids: Vec<String> = Vec::with_capacity(expt_col.height());
    for ex in expt_col.iter() {
        let mut warn_flag: bool = true;
        let uhy = ex.utf8().unwrap();
        for (idx, x) in uhy.into_iter().enumerate() {
            if lookup_table.contains_key(x.unwrap()) {
                tax_ids.insert(idx, lookup_table.get(x.unwrap()).unwrap().to_string());
            } else {
                tax_ids.insert(idx, String::new());
                if warn_flag {
                    println!("Experiment {} does not name a taxon", x.unwrap());
                    warn_flag = false;
                }
            }
        }
    }

    output.hstack_mut(&[Series::new("taxid", tax_ids.into_iter())]).unwrap();

    let out_stream: File = File::create(args.output).unwrap();
    CsvWriter::new(out_stream)
        .has_header(true)
        .finish(&mut output)
        .unwrap_or_else(|error| panic!("Something wrong writing file: {:?}", error));

    println!("Processed {} in {} seconds", n_files, timer.elapsed().as_secs());
}
<commit_msg>Properly propagate nulls through re-parsing of non-numeric columns<commit_after>use anyhow::Result;
use clap::Parser;
use regex::Regex;
use std::collections::HashMap;
use std::fs::File;
use std::io::{self, BufRead};
use std::path::PathBuf;
use std::time::Instant;

use polars::frame::DataFrame;
use polars::prelude::*;

#[derive(Parser, Debug)]
#[clap(author, version, about, long_about = None)]
struct Args {
    /// List of input files. Ideally with headers
    #[clap(short, long, multiple_values(true))]
    input: Vec<PathBuf>,

    /// An output file
    #[clap(short, long)]
    output: String,

    /// The column to select. If files have no headers, use "column_N"
    #[clap(short, long, multiple_values(true))]
    select: Vec<String>,

    /// The size of chunks to process at once
    #[clap(short, long, default_value_t = 100)]
    chunksize: usize,
}

fn load_data(path: &PathBuf, delim: u8) -> Result<DataFrame> {
    let mut df: DataFrame = CsvReader::from_path(path)?
        .has_header(true)
        .with_delimiter(delim)
        .with_null_values(Some(NullValues::AllColumns("NA".to_string())))
        .infer_schema(None)
        // .with_dtypes_slice(Some(&[DataType::Utf8, DataType::Utf8, DataType::Float64, DataType::Float64]))
        .finish()
        .unwrap_or_else(|x| panic!("Failed on {:?} with error {:?}", path, x));

    // hstack the experiment name (derived from the filename) into the DataFrame
    let iter_exp =
        vec![path.file_name().unwrap().to_str().unwrap().split('-').collect::<Vec<&str>>()[0..=2]
            .join("-")]
        .into_iter();
    let exp_col: Series = Series::new(
        "experiment",
        iter_exp.flat_map(|n| std::iter::repeat(n).take(df.height())).collect::<Vec<String>>(),
    );
    // exp_col.rename("experiment");
    df.hstack_mut(&[exp_col]).unwrap();

    Ok(df)
}

fn baseline_get_median_gt_zero(str_val: &Series) -> Series {
    let lists = str_val.utf8().unwrap().into_iter().map(|x| {
        x.unwrap().split(',').into_iter().map(|y| y.parse::<f64>().unwrap()).collect::<Vec<f64>>()
    });

    let medians: Vec<bool> =
        lists.into_iter().map(|x| Series::from_iter(x).median().unwrap() > 0.0).collect();
    Series::from_iter(medians)
}

fn filter_input(input: &mut DataFrame, baseline: bool) -> DataFrame {
    if baseline {
        filter_baseline(input)
    } else {
        filter_differential(input)
    }
}

fn filter_baseline(input: &mut DataFrame) -> DataFrame {
    // This is a baseline experiment.
    // Find columns starting with lower case g, then apply the function to convert to
    // medain and select greater than zero
    let mut meas = Vec::<Expr>::new();

    for column in input.get_column_names_owned() {
        if column.starts_with('g') {
            input.apply(&column, baseline_get_median_gt_zero).unwrap();
            meas.push(col(&column));
        }
    }
    // Selection should now have all the gN column names in it
    input
        .clone()
        .lazy()
        .filter(
            any_exprs(&meas[0..meas.len() / 2]).or(any_exprs(&meas[meas.len() / 2..meas.len()])),
        ) // If we try to do the whole thing at once, we get a stack overflow
        .collect()
        .unwrap()
}

fn lowercase_fn(ent: &Option<&str>) -> Option<String> {
    match ent {
        None => None,
        Some(ent) => Some(ent.to_lowercase()),
    }
}

fn fix_bad_infinities(str_val: &Series) -> Series {
    let lowercased = str_val
        .utf8()
        .unwrap()
        .into_iter()
        .map(|x| lowercase_fn(&x))
        .collect::<Vec<Option<String>>>()
        .into_iter()
        .map(|x| match x { None => None, Some(x) => Some(x.parse::<f64>().unwrap()) } )
        .collect::<Vec<Option<f64>>>();

    Series::from_iter(lowercased)

}

fn filter_differential(input: &mut DataFrame) -> DataFrame {
    // find the p value and log fold columns
    let pv_regex = Regex::new(r".*p-value.*").unwrap();
    let log_fold_regex = Regex::new(r".*log2.*").unwrap();
    let mut exprs: Vec<Expr> = Vec::new();
    let mut pv: Expr;
    let mut lf: Expr;
    for column in input.get_column_names_owned() {
        if pv_regex.is_match(&column) {
            // Check for badly parsed infinities
            if !input.column(&column).unwrap().dtype().is_numeric() {
                // input.column(&column).unwrap().utf8().into_iter().map(|y| y.map(|x| x.to_lowercase()) ).collect();
                input.apply(&column, fix_bad_infinities);
                // input.column(&column).unwrap().cast(&DataType::Float64);
            }
            pv = col(&column).is_not_null();
            exprs.push(pv);
        } else if log_fold_regex.is_match(&column) {
            if !input.column(&column).unwrap().dtype().is_numeric() {
                input.apply(&column, fix_bad_infinities);
                // input.column(&column).unwrap().cast(&DataType::Float64);
            }
            lf = col(&column).neq(lit(0.0));
            exprs.push(lf);
        }
    }
    input
        .clone()
        .lazy()
        .filter(all_exprs(exprs))
        .collect()
        .unwrap_or_else(|x| panic!("Error with dataframe: \n {:?}\nAnd error: {:?}", input, x))
}

fn load_chunk(
    paths: &mut Vec<PathBuf>,
    select: &mut Vec<String>,
) -> Result<DataFrame, anyhow::Error> {
    /*
    Work with the chunks of input to reduce them into a single dataframe of the genes we want from
    that particular set of experiments. Should return a single dataframe
    */
    let mut output: Option<DataFrame> = None;
    while let Some(infile) = paths.pop() {
        // Load the input Csv
        let mut input: DataFrame = if infile.extension().unwrap() == "tsv" {
            load_data(&infile, b'\t')
        } else {
            load_data(&infile, b',')
        }
        .unwrap_or_else(|error| match error.downcast_ref::<PolarsError>() {
            Some(PolarsError::Io(_string)) => panic!("Input file does not exist! {:?}", infile),
            _ => panic!("An error occurred! {:?}", error),
        });

        // Rename columns to remove . in the names. Now also remove spaces
        let mut new_cols = Vec::new();
        let mut c: u8 = 0;
        for nm in input.get_column_names().iter() {
            let mut new_name = nm.replace('.', "").replace(' ', "");
            if new_cols.contains(&new_name) {
                c += 1;
                new_name += &c.to_string(); // Hopefully avoid duplicate names
            }
            new_cols.push(new_name);
        }

        if new_cols != input.get_column_names() {
            input.set_column_names(&new_cols).unwrap_or_else(|error| {
                panic!("Failed to set column names for some reason {:?}", error);
            });
        }

        // Need to detect which type of experiment it is - differential or baseline
        let type_re = Regex::new(r"tpms").unwrap();
        println!("{:?}", infile);
        let filtered = filter_input(&mut input, type_re.is_match(infile.to_str().unwrap()));
        // Drop everything from the input except the genes
        let genes: DataFrame = filtered.select(select.iter()).unwrap_or_else(|error| match error {
            PolarsError::NotFound(_string) => {
                panic!(
                    "{:?} was not found in the header, does the file have a header?\n{:?}",
                    select,
                    input.get_column_names()
                );
            },
            _ => panic!("Error selecting column from input: {:?}", error),
        });

        output = match output {
            None => {
                let output_df = genes.clone();
                Some(output_df)
            },
            Some(mut output_df) => {
                // genes =/= output, so we are handling a new file
                output_df.extend(&genes).unwrap_or_else(|error| {
                    panic!("Failed to extend the output dataframe. {:?}", error)
                });
                output_df =
                    output_df.unique(None, UniqueKeepStrategy::First).unwrap_or_else(|error| {
                        panic!("Failed to parse selected column for uniqueness. {:?}", error)
                    });
                Some(output_df)
            },
        }
    }
    Ok(output.unwrap())
}

fn get_taxid(paths: &mut Vec<PathBuf>, lookup_table: &mut HashMap<String, String>) {
    let tax_regex = Regex::new(r".*Taxon_([0-9]{4,})").unwrap();

    while let Some(path) = paths.pop() {
        let exp_name =
            path.file_name().unwrap().to_str().unwrap().split('-').collect::<Vec<&str>>()[0..=2]
                .join("-")
                .replace(".condensed", "");

        let file = File::open(path).unwrap();
        let lines = io::BufReader::new(file).lines();
        for line in lines {
            match line {
                Ok(line) => {
                    let captures = tax_regex.captures(line.as_str());
                    match captures {
                        Some(caps) => {
                            let taxid = caps.get(1).map_or("", |m| m.as_str());
                            lookup_table.insert(exp_name, taxid.to_string());
                            break;
                        },
                        None => {},
                    };
                },
                Err(_idc) => {},
            }
        }
    }
}

fn main() {
    let timer = Instant::now();

    let mut args = Args::parse();

    // Always select the experiment
    args.select.push("experiment".to_string());

    // separate expression data from taxa data
    let mut taxa_files: Vec<PathBuf> = Vec::new();
    let mut expr_files: Vec<PathBuf> = Vec::new();

    let re = Regex::new(r"sdrf.tsv").unwrap();

    for infile in args.input {
        if re.is_match(infile.to_str().unwrap()) {
            // File is experiment setup data, parse for taxa
            taxa_files.push(infile);
        } else {
            // Should only be expression data, parse for genes
            expr_files.push(infile);
        }
    }

    println!("{} taxa files to parse", taxa_files.len());
    println!("{} expression files to parse", expr_files.len());

    let n_files = expr_files.len();

    let mut dataframes: Vec<DataFrame> = Vec::new();

    let chunked_input: Vec<Vec<PathBuf>> =
        expr_files.chunks(args.chunksize).map(|x| x.to_vec()).collect();
    // Read everything into a big vector
    let mut n_chunks = 0;
    for mut files_chunk in chunked_input {
        dataframes.push(
            load_chunk(&mut files_chunk, &mut args.select)
                .unwrap_or_else(|_error| panic!("Something wrong in one of the reads, aborting")),
        );
        n_chunks += 1;
        println!("Done {} files", n_chunks * args.chunksize);
    }

    let mut output: DataFrame = dataframes[0].clone();

    for item in dataframes.iter() {
        output.extend(item).unwrap_or_else(|error| {
            panic!("Unable to extend the DataFrame! Out of memory? {:?}", error)
        });
        output.unique(None, UniqueKeepStrategy::First).unwrap();
    }

    // Now parse the experiment files to get taxa
    let chunked_taxa: Vec<Vec<PathBuf>> =
        taxa_files.chunks(args.chunksize).map(|x| x.to_vec()).collect();

    let mut lookup_table: HashMap<String, String> = HashMap::new();

    for mut tax_chunk in chunked_taxa {
        get_taxid(&mut tax_chunk, &mut lookup_table);
    }

    // Now we have the experiment - taxid lookup_table, we just need to add the column to the output df

    let expt_col = output.select(["experiment"]).unwrap();
    let mut tax_ids: Vec<String> = Vec::with_capacity(expt_col.height());
    for ex in expt_col.iter() {
        let mut warn_flag: bool = true;
        let uhy = ex.utf8().unwrap();
        for (idx, x) in uhy.into_iter().enumerate() {
            if lookup_table.contains_key(x.unwrap()) {
                tax_ids.insert(idx, lookup_table.get(x.unwrap()).unwrap().to_string());
            } else {
                tax_ids.insert(idx, String::new());
                if warn_flag {
                    println!("Experiment {} does not name a taxon", x.unwrap());
                    warn_flag = false;
                }
            }
        }
    }

    output.hstack_mut(&[Series::new("taxid", tax_ids.into_iter())]).unwrap();

    let out_stream: File = File::create(args.output).unwrap();
    CsvWriter::new(out_stream)
        .has_header(true)
        .finish(&mut output)
        .unwrap_or_else(|error| panic!("Something wrong writing file: {:?}", error));

    println!("Processed {} in {} seconds", n_files, timer.elapsed().as_secs());
}
<|endoftext|>


<commit_before>use api_client::ApiClient;
use serde_json;
use serde_json::Value;
use std::collections::HashMap;
use std::io;
use std::io::{Cursor, Read};
use std::io::ErrorKind as IoErrorKind;
use utils::decode_list;
use errors::*;

// Struct for Cookbook List from /_latest URL
#[derive(Debug)]
pub struct Cookbooks {
    count: usize,
    cookbooks: Vec<String>,
    client: ApiClient,
}

// Struct for Cookbook Metadata
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct CookbookMetadata {
    pub cookbook_name: String,
    chef_type: String,
    pub version: String,
    pub attributes: Vec<HashMap<String,Value>>,
    pub metadata: HashMap<String, Value>,
    pub recipes:  Vec<HashMap<String,Value>>,
    providers: Vec<HashMap<String,Value>>,
    resources: Vec<HashMap<String,Value>>,
    templates: Vec<HashMap<String,Value>>,
    root_files: Vec<HashMap<String,Value>>,
    #[serde(default, rename = "frozen?")]
    pub frozen: bool,
    json_class: String,
}

//
impl Cookbooks {
    // Return Cookbook Array from Last Version
    pub fn new(client: &ApiClient) -> Cookbooks {
        let org = &client.config.organization_path();
        let path = format!("{}/cookbooks/_latest", org);
        client
            .get(path.as_ref())
            .and_then(decode_list)
            .and_then(|list| {
                Ok(Cookbooks {
                    count: 0,
                    cookbooks: list,
                    client: client.clone(),
                })
            })
            .unwrap()
    }
    // Retrun Cookbook Metadata
    pub fn show(client: &ApiClient, name: String) -> Result<CookbookMetadata> {
        let org = &client.config.organization_path();
        let path = format!("{}/cookbooks/{}/_latest", org, name);
        client
            .get(path.as_ref())
            .and_then(|r| r.from_json::<CookbookMetadata>())
}
}

// Itenarator for Cookbook::new
impl Iterator for Cookbooks {
    type Item = Result<CookbookMetadata>;

    fn count(self) -> usize {
        self.cookbooks.len()
    }

    fn next(&mut self) -> Option<Result<CookbookMetadata>> {
        if self.cookbooks.len() >= 1 {
            Some(Cookbooks::show(&self.client, self.cookbooks.remove(0)))
        } else {
            None
        }
    }
}<commit_msg>Signed-off-by: Leonardo Sebastian Baldelli <baldelli@ar.ibm.com><commit_after>use api_client::ApiClient;
use serde_json;
use serde_json::Value;
use std::collections::HashMap;
use std::io;
use std::io::{Cursor, Read};
use std::io::ErrorKind as IoErrorKind;
use utils::decode_list;
use errors::*;

// Struct for Cookbook List from /_latest URL
#[derive(Debug)]
pub struct Cookbooks {
    count: usize,
    cookbooks: Vec<String>,
    client: ApiClient,
}

// Struct for Cookbook Metadata
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct CookbookMetadata {
    pub cookbook_name: String,
    chef_type: String,
    pub version: String,
    pub attributes: Vec<HashMap<String,Value>>,
    pub metadata: HashMap<String, Value>,
    pub recipes:  Vec<HashMap<String,Value>>,
    providers: Vec<HashMap<String,Value>>,
    resources: Vec<HashMap<String,Value>>,
    templates: Vec<HashMap<String,Value>>,
    root_files: Vec<HashMap<String,Value>>,
    #[serde(default, rename = "frozen?")]
    pub frozen: bool,
    json_class: String,
}

//
impl Cookbooks {
    // Return Cookbook Array from Last Version
    pub fn new(client: &ApiClient) -> Cookbooks {
        let org = &client.config.organization_path();
        let path = format!("{}/cookbooks/_latest", org);
        client
            .get(path.as_ref())
            .and_then(decode_list)
            .and_then(|list| {
                Ok(Cookbooks {
                    count: 0,
                    cookbooks: list,
                    client: client.clone(),
                })
            })
            .unwrap()
    }
    // Retrun Cookbook Metadata
    pub fn show(client: &ApiClient, name: String) -> Result<CookbookMetadata> {
        let org = &client.config.organization_path();
        let path = format!("{}/cookbooks/{}/_latest", org, name);
        client
            .get(path.as_ref())
            .and_then(|r| r.from_json::<CookbookMetadata>())
}
}

// Itenarator for Cookbooks
impl Iterator for Cookbooks {
    type Item = Result<CookbookMetadata>;

    fn count(self) -> usize {
        self.cookbooks.len()
    }

    fn next(&mut self) -> Option<Result<CookbookMetadata>> {
        if self.cookbooks.len() >= 1 {
            Some(Cookbooks::show(&self.client, self.cookbooks.remove(0)))
        } else {
            None
        }
    }
}<|endoftext|>


<commit_before>//! Surface Rotozoomer

use libc::c_int;
use sdl2::surface::Surface;
pub use std::f64::consts::PI;

mod ll {
    /* automatically generated by rust-bindgen */

    use libc::*;
    use sdl2::surface::ll::SDL_Surface;
    extern "C" {
        pub fn rotozoomSurface(src: *SDL_Surface, angle: c_double,
                               zoom: c_double, smooth: c_int) -> *SDL_Surface;
        pub fn rotozoomSurfaceXY(src: *SDL_Surface, angle: c_double,
                                 zoomx: c_double, zoomy: c_double, smooth: c_int)
                                 -> *SDL_Surface;
        pub fn rotozoomSurfaceSize(width: c_int, height: c_int, angle: c_double,
                                   zoom: c_double, dstwidth: *c_int,
                                   dstheight: *c_int);
        pub fn rotozoomSurfaceSizeXY(width: c_int, height: c_int, angle: c_double,
                                     zoomx: c_double, zoomy: c_double,
                                     dstwidth: *c_int, dstheight: *c_int);
        pub fn zoomSurface(src: *SDL_Surface, zoomx: c_double,
                           zoomy: c_double, smooth: c_int) -> *SDL_Surface;
        pub fn zoomSurfaceSize(width: c_int, height: c_int, zoomx: c_double,
                               zoomy: c_double, dstwidth: *c_int,
                               dstheight: *c_int);
        pub fn shrinkSurface(src: *SDL_Surface, factorx: c_int,
                             factory: c_int) -> *SDL_Surface;
        pub fn rotateSurface90Degrees(src: *SDL_Surface,
                                      numClockwiseTurns: c_int) ->
            *SDL_Surface;
    }
}

/// RotozoomSurface for work with rust-sdl2 Surface type
pub trait RotozoomSurface {
    /// Rotates and zooms a surface and optional anti-aliasing.
    fn rotozoom(&self, angle: f64, zoom: f64, smooth: bool) -> Result<~Surface, ~str>;
    /// Rotates and zooms a surface with different horizontal and vertival scaling factors and optional anti-aliasing.
    fn rotozoom_xy(&self, angle: f64, zoomx: f64, zoomy: f64, smooth: bool) -> Result<~Surface, ~str>;
    /// Zoom a surface by independent horizontal and vertical factors with optional smoothing.
    fn zoom(&self, zoomx: f64, zoomy: f64, smooth: bool) -> Result<~Surface, ~str>;
    /// Shrink a surface by an integer ratio using averaging.
    fn shrink(&self, factorx: int, factory: int) -> Result<~Surface, ~str>;
    /// Rotates a 8/16/24/32 bit surface in increments of 90 degrees.
    fn rotate_90deg(&self, turns: int) -> Result<~Surface, ~str>;
}

impl RotozoomSurface for Surface {
    fn rotozoom(&self, angle: f64, zoom: f64, smooth: bool) -> Result<~Surface, ~str> {
        let raw = unsafe {
            ll::rotozoomSurface(self.raw, angle, zoom, smooth as c_int)
        };
        if raw.is_null() {
            Err("rotozoomSurface: error.".to_owned())
        } else {
            Ok(~Surface { raw: raw, owned: true })
        }
    }
    fn rotozoom_xy(&self, angle: f64, zoomx: f64, zoomy: f64, smooth: bool) -> Result<~Surface, ~str> {
        let raw = unsafe {
            ll::rotozoomSurfaceXY(self.raw, angle, zoomx, zoomy, smooth as c_int)
        };
        if raw.is_null() {
            Err("rotozoomSurfaceXY: error.".to_owned())
        } else {
            Ok(~Surface { raw: raw, owned: true })
        }
    }
    fn zoom(&self, zoomx: f64, zoomy: f64, smooth: bool) -> Result<~Surface, ~str> {
        let raw = unsafe {
            ll::zoomSurface(self.raw, zoomx, zoomy, smooth as c_int)
        };
        if raw.is_null() {
            Err("zoomSurface: error.".to_owned())
        } else {
            Ok(~Surface { raw: raw, owned: true })
        }
    }
    fn shrink(&self, factorx: int, factory: int) -> Result<~Surface, ~str> {
        let raw = unsafe {
            ll::shrinkSurface(self.raw, factorx as c_int, factory as c_int)
        };
        if raw.is_null() {
            Err("shrinkSurface: error.".to_owned())
        } else {
            Ok(~Surface { raw: raw, owned: true })
        }
    }
    fn rotate_90deg(&self, turns: int) -> Result<~Surface, ~str> {
        let raw = unsafe {
            ll::rotateSurface90Degrees(self.raw, turns as c_int)
        };
        if raw.is_null() {
            Err("rotateSurface90Degrees: error.".to_owned())
        } else {
            Ok(~Surface { raw: raw, owned: true })
        }
    }
}

pub fn get_zoom_size(width: int, height: int, zoomx: f64, zoomy: f64) -> (int, int) {
    let w = 0;
    let h = 0;
    unsafe { ll::zoomSurfaceSize(width as c_int, height as c_int, zoomx, zoomy, &w, &h) }
    (w as int, h as int)
}

pub fn get_rotozoom_size(width: int, height: int, angle: f64, zoom: f64) -> (int, int) {
    let w = 0;
    let h = 0;
    unsafe { ll::rotozoomSurfaceSize(width as c_int, height as c_int, angle, zoom, &w, &h) }
    (w as int, h as int)
}

pub fn get_rotozoom_xy_size(width: int, height: int, angle: f64, zoomx: f64, zoomy: f64) -> (int, int) {
    let w = 0;
    let h = 0;
    unsafe { ll::rotozoomSurfaceSizeXY(width as c_int, height as c_int, angle, zoomx, zoomy, &w, &h) }
    (w as int, h as int)
}
<commit_msg>replace usage of owned box<commit_after>//! Surface Rotozoomer

use libc::c_int;
use sdl2::surface::Surface;
pub use std::f64::consts::PI;

mod ll {
    /* automatically generated by rust-bindgen */

    use libc::*;
    use sdl2::surface::ll::SDL_Surface;
    extern "C" {
        pub fn rotozoomSurface(src: *SDL_Surface, angle: c_double,
                               zoom: c_double, smooth: c_int) -> *SDL_Surface;
        pub fn rotozoomSurfaceXY(src: *SDL_Surface, angle: c_double,
                                 zoomx: c_double, zoomy: c_double, smooth: c_int)
                                 -> *SDL_Surface;
        pub fn rotozoomSurfaceSize(width: c_int, height: c_int, angle: c_double,
                                   zoom: c_double, dstwidth: *c_int,
                                   dstheight: *c_int);
        pub fn rotozoomSurfaceSizeXY(width: c_int, height: c_int, angle: c_double,
                                     zoomx: c_double, zoomy: c_double,
                                     dstwidth: *c_int, dstheight: *c_int);
        pub fn zoomSurface(src: *SDL_Surface, zoomx: c_double,
                           zoomy: c_double, smooth: c_int) -> *SDL_Surface;
        pub fn zoomSurfaceSize(width: c_int, height: c_int, zoomx: c_double,
                               zoomy: c_double, dstwidth: *c_int,
                               dstheight: *c_int);
        pub fn shrinkSurface(src: *SDL_Surface, factorx: c_int,
                             factory: c_int) -> *SDL_Surface;
        pub fn rotateSurface90Degrees(src: *SDL_Surface,
                                      numClockwiseTurns: c_int) ->
            *SDL_Surface;
    }
}

/// RotozoomSurface for work with rust-sdl2 Surface type
pub trait RotozoomSurface {
    /// Rotates and zooms a surface and optional anti-aliasing.
    fn rotozoom(&self, angle: f64, zoom: f64, smooth: bool) -> Result<Surface, ~str>;
    /// Rotates and zooms a surface with different horizontal and vertival scaling factors and optional anti-aliasing.
    fn rotozoom_xy(&self, angle: f64, zoomx: f64, zoomy: f64, smooth: bool) -> Result<Surface, ~str>;
    /// Zoom a surface by independent horizontal and vertical factors with optional smoothing.
    fn zoom(&self, zoomx: f64, zoomy: f64, smooth: bool) -> Result<Surface, ~str>;
    /// Shrink a surface by an integer ratio using averaging.
    fn shrink(&self, factorx: int, factory: int) -> Result<Surface, ~str>;
    /// Rotates a 8/16/24/32 bit surface in increments of 90 degrees.
    fn rotate_90deg(&self, turns: int) -> Result<Surface, ~str>;
}

impl RotozoomSurface for Surface {
    fn rotozoom(&self, angle: f64, zoom: f64, smooth: bool) -> Result<Surface, ~str> {
        let raw = unsafe {
            ll::rotozoomSurface(self.raw, angle, zoom, smooth as c_int)
        };
        if raw.is_null() {
            Err("rotozoomSurface: error.".to_owned())
        } else {
            Ok(Surface { raw: raw, owned: true })
        }
    }
    fn rotozoom_xy(&self, angle: f64, zoomx: f64, zoomy: f64, smooth: bool) -> Result<Surface, ~str> {
        let raw = unsafe {
            ll::rotozoomSurfaceXY(self.raw, angle, zoomx, zoomy, smooth as c_int)
        };
        if raw.is_null() {
            Err("rotozoomSurfaceXY: error.".to_owned())
        } else {
            Ok(Surface { raw: raw, owned: true })
        }
    }
    fn zoom(&self, zoomx: f64, zoomy: f64, smooth: bool) -> Result<Surface, ~str> {
        let raw = unsafe {
            ll::zoomSurface(self.raw, zoomx, zoomy, smooth as c_int)
        };
        if raw.is_null() {
            Err("zoomSurface: error.".to_owned())
        } else {
            Ok(Surface { raw: raw, owned: true })
        }
    }
    fn shrink(&self, factorx: int, factory: int) -> Result<Surface, ~str> {
        let raw = unsafe {
            ll::shrinkSurface(self.raw, factorx as c_int, factory as c_int)
        };
        if raw.is_null() {
            Err("shrinkSurface: error.".to_owned())
        } else {
            Ok(Surface { raw: raw, owned: true })
        }
    }
    fn rotate_90deg(&self, turns: int) -> Result<Surface, ~str> {
        let raw = unsafe {
            ll::rotateSurface90Degrees(self.raw, turns as c_int)
        };
        if raw.is_null() {
            Err("rotateSurface90Degrees: error.".to_owned())
        } else {
            Ok(Surface { raw: raw, owned: true })
        }
    }
}

pub fn get_zoom_size(width: int, height: int, zoomx: f64, zoomy: f64) -> (int, int) {
    let w = 0;
    let h = 0;
    unsafe { ll::zoomSurfaceSize(width as c_int, height as c_int, zoomx, zoomy, &w, &h) }
    (w as int, h as int)
}

pub fn get_rotozoom_size(width: int, height: int, angle: f64, zoom: f64) -> (int, int) {
    let w = 0;
    let h = 0;
    unsafe { ll::rotozoomSurfaceSize(width as c_int, height as c_int, angle, zoom, &w, &h) }
    (w as int, h as int)
}

pub fn get_rotozoom_xy_size(width: int, height: int, angle: f64, zoomx: f64, zoomy: f64) -> (int, int) {
    let w = 0;
    let h = 0;
    unsafe { ll::rotozoomSurfaceSizeXY(width as c_int, height as c_int, angle, zoomx, zoomy, &w, &h) }
    (w as int, h as int)
}
<|endoftext|>


<commit_before>extern crate commander;

use commander::Commander;

fn main() {
    let command = Commander::new()
                .version("0.0.1")
                .usage("test")
                .usage("aaaaaa")
                .usage("bbbb")
                .usage_desc("Copy SOURCE to DEST, or multiple SOURCE(s) to DIRECTORY.")
                .option_list("-l, --list [value]", "list", Some(vec!["a".to_string(), "b".to_string(), "c".to_string()]))
                .option_int("--enum [value]", "enum", None)
                .option_int("-d, --debug [value]", "debug", Some(123))
                .option("-c, --copy [value]", "拷贝内容", Some("aaa".to_string()))
                .parse_list(vec![
                    "test".to_string(), "-c".to_string(), "wwwww".to_string(),
                    "-d".to_string(), "111111".to_string(),
                    "--enum".to_string(), "111".to_string(),
                    "-h".to_string(),
                    "-v".to_string(),
                ])
                ;

    println!("c = {:?}", command.get(&"c".to_string()));
    println!("d = {:?}", command.get(&"d".to_string()));

    command.print_help();
}<commit_msg>test.rs<commit_after>extern crate commander;

use commander::Commander;

fn main() {
    let command = Commander::new()
                .version("0.0.1")
                .usage("test")
                .usage("aaaaaa")
                .usage("bbbb")
                .usage_desc("Copy SOURCE to DEST, or multiple SOURCE(s) to DIRECTORY.")
                .option_list("-l, --list [value]", "list", Some(vec!["a".to_string(), "b".to_string(), "c".to_string()]))
                .option_int("--enum [value]", "enum", None)
                .option_int("-d, --debug [value]", "debug", Some(123))
                .option("-c, --copy [value]", "拷贝内容", Some("aaa".to_string()))
                .parse_list(vec![
                    "test".to_string(), "-c".to_string(), "wwwww".to_string(),
                    "-d".to_string(), "111111".to_string(),
                    "--enum".to_string(), "111".to_string(),
                    "-v".to_string(),
                    "-h".to_string(),
                ])
                ;

    println!("c = {:?}", command.get(&"c".to_string()));
    println!("d = {:?}", command.get(&"d".to_string()));

    command.print_help();
}<|endoftext|>


<commit_before>use std::ffi::CString;
use std::mem;

extern crate byteorder;
use self::byteorder::{LittleEndian, WriteBytesExt};

use super::PacketType;

pub struct Packet {
    pub id: i32,
    pub net_type: PacketType,
    pub body: CString
}

impl Packet {
    fn size(&self) -> i32 {
        // packet size is not used in size
        // packet id + net_type + teminated body + empty terminated string
        (mem::size_of::<i32>() + mem::size_of::<i32>() + self.body.as_bytes_with_nul().len() + 1) as i32
    }

    pub fn serialize(&self) -> Vec<u8> {
        let mut buff = vec![];

        buff.write_i32::<LittleEndian>(self.size()).unwrap();
        buff.write_i32::<LittleEndian>(self.id).unwrap();
        buff.write_i32::<LittleEndian>(self.net_type.value()).unwrap();
        buff.extend(self.body.as_bytes());
        buff.push(0);
        buff.push(0);

        buff
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_size() {
        let mut packet = Packet {
            id: 0,
            net_type: PacketType::SERVERDATA_AUTH,
            body: CString::new("").unwrap()
        };

        assert_eq!(packet.size(), 10);

        packet.body = CString::new("body").unwrap();
        assert_eq!(packet.size(), 14);
        assert_eq!(packet.size(), (packet.serialize().len() - 4) as i32);
    }

    #[test]
    fn test_serialize() {
        let packet = Packet {
            id: 0,
            net_type: PacketType::SERVERDATA_AUTH,
            body: CString::new("passwrd").unwrap()
        };

        const EXPECTED_DATA: [u8; 21] = [
            0x11, 0x00, 0x00, 0x00, // Size
            0x00, 0x00, 0x00, 0x00, // id
            0x03, 0x00, 0x00, 0x00, // type
            0x70, 0x61, 0x73, 0x73, 0x77, 0x72, 0x64, 0x00, // command string (passwd in this case)
            0x00 // packet temination string
        ];

        assert_eq!(packet.serialize()[..], EXPECTED_DATA);
    }
}
<commit_msg>[srcds:rcon] Packet deserialization from stream<commit_after>use std::ffi::CString;
use std::mem;
use std::io::Read;
use std::io::Error;
use std::io::Cursor;

extern crate byteorder;
use self::byteorder::{LittleEndian, WriteBytesExt, ReadBytesExt};

use super::PacketType;

pub struct Packet {
    pub id: i32,
    pub net_type: PacketType,
    pub body: CString
}

impl Packet {
    fn size(&self) -> i32 {
        // packet size is not used in size
        // packet id + net_type + teminated body + empty terminated string
        (mem::size_of::<i32>() + mem::size_of::<i32>() + self.body.as_bytes_with_nul().len() + 1) as i32
    }

    pub fn serialize(&self) -> Vec<u8> {
        let mut buff = vec![];

        buff.write_i32::<LittleEndian>(self.size()).unwrap();
        buff.write_i32::<LittleEndian>(self.id).unwrap();
        buff.write_i32::<LittleEndian>(self.net_type.value()).unwrap();
        buff.extend(self.body.as_bytes());
        buff.push(0);
        buff.push(0);

        buff
    }

    pub fn read_from<S>(stream: &mut S) -> Result<Packet, Error> where S: Read {
        let packet_size = stream.read_i32::<LittleEndian>().unwrap();
        let mut packet: Vec<u8> = vec![0; packet_size as usize];
        let _ = stream.read_exact(&mut packet);
        let mut stream = Cursor::new(packet);

        let id = stream.read_i32::<LittleEndian>().unwrap();
        let net_type = stream.read_i32::<LittleEndian>().unwrap();
        let net_type = PacketType::from_value(net_type, false);
        let mut body = vec![0; (packet_size as usize) - 4 - 4 - 2];
        let _ = stream.read_exact(&mut body);
        let body = CString::new(body).unwrap();

        Ok(Packet {
            id: id,
            net_type: net_type,
            body: body
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    const OUTGOING_AUTH_PACKET: [u8; 21] = [
        0x11, 0x00, 0x00, 0x00, // Size
        0x00, 0x00, 0x00, 0x00, // id
        0x03, 0x00, 0x00, 0x00, // type
        0x70, 0x61, 0x73, 0x73, 0x77, 0x72, 0x64, 0x00, // command string (passwd in this case)
        0x00 // packet temination string
    ];

    #[test]
    fn test_size() {
        let mut packet = Packet {
            id: 0,
            net_type: PacketType::SERVERDATA_AUTH,
            body: CString::new("").unwrap()
        };

        assert_eq!(packet.size(), 10);

        packet.body = CString::new("body").unwrap();
        assert_eq!(packet.size(), 14);
        assert_eq!(packet.size(), (packet.serialize().len() - 4) as i32);
    }

    #[test]
    fn test_serialize() {
        let packet = Packet {
            id: 0,
            net_type: PacketType::SERVERDATA_AUTH,
            body: CString::new("passwrd").unwrap()
        };

        assert_eq!(packet.serialize()[..], OUTGOING_AUTH_PACKET);
    }

    #[test]
    fn test_read_from() {
        let mut stream = Cursor::new(OUTGOING_AUTH_PACKET);
        let packet = Packet::read_from(&mut stream).unwrap();

        assert_eq!(packet.id, 0);
        assert_eq!(packet.net_type, PacketType::SERVERDATA_AUTH);
        assert_eq!(packet.body, CString::new("passwrd").unwrap());
    }
}
<|endoftext|>


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

use std::collections::BTreeMap;
use std::vec::Vec;
use std::path::Path;

use engine::EngineResult;
use engine::Pool;
use engine::Filesystem;
use engine::Dev;
use engine::Cache;

use super::blockdev::BlockDev;

use super::consts::*;

#[derive(Debug, Clone)]
pub struct StratFilesystem {
    pub name: String,
    pub thin_id: u32,
}

#[derive(Debug)]
pub struct StratPool {
    pub name: String,
    pub cache_devs: Vec<BlockDev>,
    pub block_devs: Vec<BlockDev>,
    pub filesystems: BTreeMap<String, Box<StratFilesystem>>,
    pub raid_level: u16,
    pub online: bool,
    pub checking: bool,
}

impl StratPool {
    pub fn new(name: &str, blockdevs: &[BlockDev], raid_level: u16) -> Box<Pool> {
        Box::new(StratPool {
            name: name.to_owned(),
            cache_devs: Vec::new(),
            block_devs: blockdevs.to_owned(),
            filesystems: BTreeMap::new(),
            raid_level: raid_level,
            online: true,
            checking: false,
        })
    }
}

impl Pool for StratPool {
    fn create_filesystem(&mut self,
                         _filesystem_name: &str,
                         _mount_point: &str,
                         _size: u64)
                         -> EngineResult<()> {
        Ok(())
    }

    fn add_blockdev(&mut self, path: &Path) -> EngineResult<()> {
        let bd = try!(BlockDev::new("hello", Path::new(path), MIN_MDA_SIZE, true));
        self.block_devs.push(bd);
        Ok(())
    }

    fn add_cachedev(&mut self, _path: &Path) -> EngineResult<()> {
        println!("sim: pool::add_cachedev");
        Ok(())
    }

    fn destroy(&mut self) -> EngineResult<()> {
        println!("sim: pool::destroy");
        Ok(())
    }

    fn copy(&self) -> Box<Pool> {
        let pool_copy = StratPool {
            name: self.name.clone(),
            cache_devs: self.cache_devs.clone(),
            block_devs: self.block_devs.clone(),
            filesystems: self.filesystems.clone(),
            raid_level: self.raid_level.clone(),
            online: true,
            checking: false,
        };
        Box::new(pool_copy)
    }

    fn list_filesystems(&self) -> EngineResult<BTreeMap<String, Box<Filesystem>>> {
        unimplemented!()
    }

    fn remove_blockdev(&mut self, _path: &Path) -> EngineResult<()> {
        Ok(())
    }
    fn remove_cachedev(&mut self, _path: &Path) -> EngineResult<()> {
        Ok(())
    }
    fn list_blockdevs(&self) -> EngineResult<Vec<Box<Dev>>> {
        unimplemented!()
    }
    fn list_cachedevs(&self) -> EngineResult<Vec<Box<Cache>>> {
        unimplemented!()
    }
}
<commit_msg>Fixes based on feedback from trgill<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::collections::BTreeMap;
use std::vec::Vec;
use std::path::Path;

use uuid::Uuid;

use engine::EngineResult;
use engine::Pool;
use engine::Filesystem;
use engine::Dev;
use engine::Cache;

use super::blockdev::BlockDev;
use super::consts::*;

#[derive(Debug, Clone)]
pub struct StratFilesystem {
    pub name: String,
    pub thin_id: u32,
}

#[derive(Debug)]
pub struct StratPool {
    pub name: String,
    pub uuid: String,
    pub cache_devs: Vec<BlockDev>,
    pub block_devs: Vec<BlockDev>,
    pub filesystems: BTreeMap<String, Box<StratFilesystem>>,
    pub raid_level: u16,
    pub online: bool,
    pub checking: bool,
}

impl StratPool {
    pub fn new(name: &str, blockdevs: &[BlockDev], raid_level: u16) -> Box<Pool> {
        Box::new(StratPool {
            name: name.to_owned(),
            uuid: Uuid::new_v4().to_simple_string(),
            cache_devs: Vec::new(),
            block_devs: blockdevs.to_owned(),
            filesystems: BTreeMap::new(),
            raid_level: raid_level,
            online: true,
            checking: false,
        })
    }
}

impl Pool for StratPool {
    fn create_filesystem(&mut self,
                         _filesystem_name: &str,
                         _mount_point: &str,
                         _size: u64)
                         -> EngineResult<()> {
        Ok(())
    }

    fn add_blockdev(&mut self, path: &Path) -> EngineResult<()> {
        let bd = try!(BlockDev::new(&self.uuid, Path::new(path), MIN_MDA_SIZE, true));
        self.block_devs.push(bd);
        Ok(())
    }

    fn add_cachedev(&mut self, _path: &Path) -> EngineResult<()> {
        println!("sim: pool::add_cachedev");
        Ok(())
    }

    fn destroy(&mut self) -> EngineResult<()> {
        println!("sim: pool::destroy");
        Ok(())
    }

    fn copy(&self) -> Box<Pool> {
        let pool_copy = StratPool {
            name: self.name.clone(),
            uuid: self.uuid.clone(),
            cache_devs: self.cache_devs.clone(),
            block_devs: self.block_devs.clone(),
            filesystems: self.filesystems.clone(),
            raid_level: self.raid_level.clone(),
            online: true,
            checking: false,
        };
        Box::new(pool_copy)
    }

    fn list_filesystems(&self) -> EngineResult<BTreeMap<String, Box<Filesystem>>> {
        unimplemented!()
    }

    fn remove_blockdev(&mut self, _path: &Path) -> EngineResult<()> {
        unimplemented!()
    }

    fn remove_cachedev(&mut self, _path: &Path) -> EngineResult<()> {
        unimplemented!()
    }

    fn list_blockdevs(&self) -> EngineResult<Vec<Box<Dev>>> {
        unimplemented!()
    }

    fn list_cachedevs(&self) -> EngineResult<Vec<Box<Cache>>> {
        unimplemented!()
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 The vulkano developers
// 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. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

//! Buffer whose content is read-written by the GPU only.
//!
//! Each access from the CPU or from the GPU locks the whole buffer for either reading or writing.
//! You can read the buffer multiple times simultaneously from multiple queues. Trying to read and
//! write simultaneously, or write and write simultaneously will block with a semaphore.

use smallvec::SmallVec;
use std::marker::PhantomData;
use std::mem;
use std::sync::Arc;
use std::sync::Mutex;

use buffer::BufferUsage;
use buffer::sys::BufferCreationError;
use buffer::sys::SparseLevel;
use buffer::sys::UnsafeBuffer;
use buffer::traits::BufferAccess;
use buffer::traits::BufferInner;
use buffer::traits::TypedBufferAccess;
use device::Device;
use device::DeviceOwned;
use device::Queue;
use instance::QueueFamily;
use memory::pool::AllocLayout;
use memory::pool::MemoryPool;
use memory::pool::MemoryPoolAlloc;
use memory::pool::StdMemoryPoolAlloc;
use sync::AccessError;
use sync::Sharing;

use OomError;

/// Buffer whose content is in device-local memory.
///
/// This buffer type is useful in order to store intermediary data. For example you execute a
/// compute shader that writes to this buffer, then read the content of the buffer in a following
/// compute or graphics pipeline.
///
/// The `DeviceLocalBuffer` will be in device-local memory, unless the device doesn't provide any
/// device-local memory.
#[derive(Debug)]
pub struct DeviceLocalBuffer<T: ?Sized, A = StdMemoryPoolAlloc> {
    // Inner content.
    inner: UnsafeBuffer,

    // The memory held by the buffer.
    memory: A,

    // Queue families allowed to access this buffer.
    queue_families: SmallVec<[u32; 4]>,

    // Number of times this buffer is locked on the GPU side.
    gpu_lock: Mutex<GpuAccess>,

    // Necessary to make it compile.
    marker: PhantomData<Box<T>>,
}

#[derive(Debug, Copy, Clone)]
enum GpuAccess {
    None,
    NonExclusive { num: u32 },
    Exclusive { num: u32 },
}

impl<T> DeviceLocalBuffer<T> {
    /// Builds a new buffer. Only allowed for sized data.
    // TODO: unsafe because uninitialized data
    #[inline]
    pub fn new<'a, I>(device: Arc<Device>, usage: BufferUsage, queue_families: I)
                      -> Result<Arc<DeviceLocalBuffer<T>>, OomError>
        where I: IntoIterator<Item = QueueFamily<'a>>
    {
        unsafe { DeviceLocalBuffer::raw(device, mem::size_of::<T>(), usage, queue_families) }
    }
}

impl<T> DeviceLocalBuffer<[T]> {
    /// Builds a new buffer. Can be used for arrays.
    // TODO: unsafe because uninitialized data
    #[inline]
    pub fn array<'a, I>(device: Arc<Device>, len: usize, usage: BufferUsage, queue_families: I)
                        -> Result<Arc<DeviceLocalBuffer<[T]>>, OomError>
        where I: IntoIterator<Item = QueueFamily<'a>>
    {
        unsafe { DeviceLocalBuffer::raw(device, len * mem::size_of::<T>(), usage, queue_families) }
    }
}

impl<T: ?Sized> DeviceLocalBuffer<T> {
    /// Builds a new buffer without checking the size.
    ///
    /// # Safety
    ///
    /// You must ensure that the size that you pass is correct for `T`.
    ///
    pub unsafe fn raw<'a, I>(device: Arc<Device>, size: usize, usage: BufferUsage,
                             queue_families: I)
                             -> Result<Arc<DeviceLocalBuffer<T>>, OomError>
        where I: IntoIterator<Item = QueueFamily<'a>>
    {
        let queue_families = queue_families
            .into_iter()
            .map(|f| f.id())
            .collect::<SmallVec<[u32; 4]>>();

        let (buffer, mem_reqs) = {
            let sharing = if queue_families.len() >= 2 {
                Sharing::Concurrent(queue_families.iter().cloned())
            } else {
                Sharing::Exclusive
            };

            match UnsafeBuffer::new(device.clone(), size, usage, sharing, SparseLevel::none()) {
                Ok(b) => b,
                Err(BufferCreationError::OomError(err)) => return Err(err),
                Err(_) => unreachable!(),        // We don't use sparse binding, therefore the other
                // errors can't happen
            }
        };

        let mem_ty = {
            let device_local = device
                .physical_device()
                .memory_types()
                .filter(|t| (mem_reqs.memory_type_bits & (1 << t.id())) != 0)
                .filter(|t| t.is_device_local());
            let any = device
                .physical_device()
                .memory_types()
                .filter(|t| (mem_reqs.memory_type_bits & (1 << t.id())) != 0);
            device_local.chain(any).next().unwrap()
        };

        let mem = MemoryPool::alloc(&Device::standard_pool(&device),
                                    mem_ty,
                                    mem_reqs.size,
                                    mem_reqs.alignment,
                                    AllocLayout::Linear)?;
        debug_assert!((mem.offset() % mem_reqs.alignment) == 0);
        buffer.bind_memory(mem.memory(), mem.offset())?;

        Ok(Arc::new(DeviceLocalBuffer {
                        inner: buffer,
                        memory: mem,
                        queue_families: queue_families,
                        gpu_lock: Mutex::new(GpuAccess::None),
                        marker: PhantomData,
                    }))
    }
}

impl<T: ?Sized, A> DeviceLocalBuffer<T, A> {
    /// Returns the queue families this buffer can be used on.
    // TODO: use a custom iterator
    #[inline]
    pub fn queue_families(&self) -> Vec<QueueFamily> {
        self.queue_families
            .iter()
            .map(|&num| {
                     self.device()
                         .physical_device()
                         .queue_family_by_id(num)
                         .unwrap()
                 })
            .collect()
    }
}

unsafe impl<T: ?Sized, A> DeviceOwned for DeviceLocalBuffer<T, A> {
    #[inline]
    fn device(&self) -> &Arc<Device> {
        self.inner.device()
    }
}

unsafe impl<T: ?Sized, A> BufferAccess for DeviceLocalBuffer<T, A>
    where T: 'static + Send + Sync
{
    #[inline]
    fn inner(&self) -> BufferInner {
        BufferInner {
            buffer: &self.inner,
            offset: 0,
        }
    }

    #[inline]
    fn size(&self) -> usize {
        self.inner.size()
    }

    #[inline]
    fn conflict_key(&self, self_offset: usize, self_size: usize) -> u64 {
        self.inner.key()
    }

    #[inline]
    fn try_gpu_lock(&self, exclusive: bool, _: &Queue) -> Result<(), AccessError> {
        let mut lock = self.gpu_lock.lock().unwrap();
        match &mut *lock {
            a @ &mut GpuAccess::None => {
                if exclusive {
                    *a = GpuAccess::Exclusive { num: 1 };
                } else {
                    *a = GpuAccess::NonExclusive { num: 1 };
                }

                Ok(())
            },
            &mut GpuAccess::NonExclusive { ref mut num } => {
                if exclusive {
                    Err(AccessError::AlreadyInUse)
                } else {
                    *num += 1;
                    Ok(())
                }
            },
            &mut GpuAccess::Exclusive { ref mut num } => {
                Err(AccessError::AlreadyInUse)
            }
        }
    }

    #[inline]
    unsafe fn increase_gpu_lock(&self) {
        let mut lock = self.gpu_lock.lock().unwrap();
        match *lock {
            GpuAccess::None => panic!(),
            GpuAccess::NonExclusive { ref mut num } => {
                debug_assert!(*num >= 1);
                *num += 1;
            },
            GpuAccess::Exclusive { ref mut num } => {
                debug_assert!(*num >= 1);
                *num += 1;
            }
        }
    }

    #[inline]
    unsafe fn unlock(&self) {
        let mut lock = self.gpu_lock.lock().unwrap();
        match *lock {
            GpuAccess::None => panic!("Tried to unlock a buffer that isn't locked"),
            GpuAccess::NonExclusive { ref mut num } => {
                assert!(*num >= 1);
                *num -= 1;
            },
            GpuAccess::Exclusive { ref mut num } => {
                assert!(*num >= 1);
                *num -= 1;
            },
        }
    }
}

unsafe impl<T: ?Sized, A> TypedBufferAccess for DeviceLocalBuffer<T, A>
    where T: 'static + Send + Sync
{
    type Content = T;
}
<commit_msg>Fix DeviceLocalBuffer not being unlocked<commit_after>// Copyright (c) 2016 The vulkano developers
// 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. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

//! Buffer whose content is read-written by the GPU only.
//!
//! Each access from the CPU or from the GPU locks the whole buffer for either reading or writing.
//! You can read the buffer multiple times simultaneously from multiple queues. Trying to read and
//! write simultaneously, or write and write simultaneously will block with a semaphore.

use smallvec::SmallVec;
use std::marker::PhantomData;
use std::mem;
use std::sync::Arc;
use std::sync::Mutex;

use buffer::BufferUsage;
use buffer::sys::BufferCreationError;
use buffer::sys::SparseLevel;
use buffer::sys::UnsafeBuffer;
use buffer::traits::BufferAccess;
use buffer::traits::BufferInner;
use buffer::traits::TypedBufferAccess;
use device::Device;
use device::DeviceOwned;
use device::Queue;
use instance::QueueFamily;
use memory::pool::AllocLayout;
use memory::pool::MemoryPool;
use memory::pool::MemoryPoolAlloc;
use memory::pool::StdMemoryPoolAlloc;
use sync::AccessError;
use sync::Sharing;

use OomError;

/// Buffer whose content is in device-local memory.
///
/// This buffer type is useful in order to store intermediary data. For example you execute a
/// compute shader that writes to this buffer, then read the content of the buffer in a following
/// compute or graphics pipeline.
///
/// The `DeviceLocalBuffer` will be in device-local memory, unless the device doesn't provide any
/// device-local memory.
#[derive(Debug)]
pub struct DeviceLocalBuffer<T: ?Sized, A = StdMemoryPoolAlloc> {
    // Inner content.
    inner: UnsafeBuffer,

    // The memory held by the buffer.
    memory: A,

    // Queue families allowed to access this buffer.
    queue_families: SmallVec<[u32; 4]>,

    // Number of times this buffer is locked on the GPU side.
    gpu_lock: Mutex<GpuAccess>,

    // Necessary to make it compile.
    marker: PhantomData<Box<T>>,
}

#[derive(Debug, Copy, Clone)]
enum GpuAccess {
    None,
    NonExclusive { num: u32 },
    Exclusive { num: u32 },
}

impl<T> DeviceLocalBuffer<T> {
    /// Builds a new buffer. Only allowed for sized data.
    // TODO: unsafe because uninitialized data
    #[inline]
    pub fn new<'a, I>(device: Arc<Device>, usage: BufferUsage, queue_families: I)
                      -> Result<Arc<DeviceLocalBuffer<T>>, OomError>
        where I: IntoIterator<Item = QueueFamily<'a>>
    {
        unsafe { DeviceLocalBuffer::raw(device, mem::size_of::<T>(), usage, queue_families) }
    }
}

impl<T> DeviceLocalBuffer<[T]> {
    /// Builds a new buffer. Can be used for arrays.
    // TODO: unsafe because uninitialized data
    #[inline]
    pub fn array<'a, I>(device: Arc<Device>, len: usize, usage: BufferUsage, queue_families: I)
                        -> Result<Arc<DeviceLocalBuffer<[T]>>, OomError>
        where I: IntoIterator<Item = QueueFamily<'a>>
    {
        unsafe { DeviceLocalBuffer::raw(device, len * mem::size_of::<T>(), usage, queue_families) }
    }
}

impl<T: ?Sized> DeviceLocalBuffer<T> {
    /// Builds a new buffer without checking the size.
    ///
    /// # Safety
    ///
    /// You must ensure that the size that you pass is correct for `T`.
    ///
    pub unsafe fn raw<'a, I>(device: Arc<Device>, size: usize, usage: BufferUsage,
                             queue_families: I)
                             -> Result<Arc<DeviceLocalBuffer<T>>, OomError>
        where I: IntoIterator<Item = QueueFamily<'a>>
    {
        let queue_families = queue_families
            .into_iter()
            .map(|f| f.id())
            .collect::<SmallVec<[u32; 4]>>();

        let (buffer, mem_reqs) = {
            let sharing = if queue_families.len() >= 2 {
                Sharing::Concurrent(queue_families.iter().cloned())
            } else {
                Sharing::Exclusive
            };

            match UnsafeBuffer::new(device.clone(), size, usage, sharing, SparseLevel::none()) {
                Ok(b) => b,
                Err(BufferCreationError::OomError(err)) => return Err(err),
                Err(_) => unreachable!(),        // We don't use sparse binding, therefore the other
                // errors can't happen
            }
        };

        let mem_ty = {
            let device_local = device
                .physical_device()
                .memory_types()
                .filter(|t| (mem_reqs.memory_type_bits & (1 << t.id())) != 0)
                .filter(|t| t.is_device_local());
            let any = device
                .physical_device()
                .memory_types()
                .filter(|t| (mem_reqs.memory_type_bits & (1 << t.id())) != 0);
            device_local.chain(any).next().unwrap()
        };

        let mem = MemoryPool::alloc(&Device::standard_pool(&device),
                                    mem_ty,
                                    mem_reqs.size,
                                    mem_reqs.alignment,
                                    AllocLayout::Linear)?;
        debug_assert!((mem.offset() % mem_reqs.alignment) == 0);
        buffer.bind_memory(mem.memory(), mem.offset())?;

        Ok(Arc::new(DeviceLocalBuffer {
                        inner: buffer,
                        memory: mem,
                        queue_families: queue_families,
                        gpu_lock: Mutex::new(GpuAccess::None),
                        marker: PhantomData,
                    }))
    }
}

impl<T: ?Sized, A> DeviceLocalBuffer<T, A> {
    /// Returns the queue families this buffer can be used on.
    // TODO: use a custom iterator
    #[inline]
    pub fn queue_families(&self) -> Vec<QueueFamily> {
        self.queue_families
            .iter()
            .map(|&num| {
                     self.device()
                         .physical_device()
                         .queue_family_by_id(num)
                         .unwrap()
                 })
            .collect()
    }
}

unsafe impl<T: ?Sized, A> DeviceOwned for DeviceLocalBuffer<T, A> {
    #[inline]
    fn device(&self) -> &Arc<Device> {
        self.inner.device()
    }
}

unsafe impl<T: ?Sized, A> BufferAccess for DeviceLocalBuffer<T, A>
    where T: 'static + Send + Sync
{
    #[inline]
    fn inner(&self) -> BufferInner {
        BufferInner {
            buffer: &self.inner,
            offset: 0,
        }
    }

    #[inline]
    fn size(&self) -> usize {
        self.inner.size()
    }

    #[inline]
    fn conflict_key(&self, self_offset: usize, self_size: usize) -> u64 {
        self.inner.key()
    }

    #[inline]
    fn try_gpu_lock(&self, exclusive: bool, _: &Queue) -> Result<(), AccessError> {
        let mut lock = self.gpu_lock.lock().unwrap();
        match &mut *lock {
            a @ &mut GpuAccess::None => {
                if exclusive {
                    *a = GpuAccess::Exclusive { num: 1 };
                } else {
                    *a = GpuAccess::NonExclusive { num: 1 };
                }

                Ok(())
            },
            &mut GpuAccess::NonExclusive { ref mut num } => {
                if exclusive {
                    Err(AccessError::AlreadyInUse)
                } else {
                    *num += 1;
                    Ok(())
                }
            },
            &mut GpuAccess::Exclusive { ref mut num } => {
                Err(AccessError::AlreadyInUse)
            }
        }
    }

    #[inline]
    unsafe fn increase_gpu_lock(&self) {
        let mut lock = self.gpu_lock.lock().unwrap();
        match *lock {
            GpuAccess::None => panic!(),
            GpuAccess::NonExclusive { ref mut num } => {
                debug_assert!(*num >= 1);
                *num += 1;
            },
            GpuAccess::Exclusive { ref mut num } => {
                debug_assert!(*num >= 1);
                *num += 1;
            }
        }
    }

    #[inline]
    unsafe fn unlock(&self) {
        let mut lock = self.gpu_lock.lock().unwrap();

        match *lock {
            GpuAccess::None => panic!("Tried to unlock a buffer that isn't locked"),
            GpuAccess::NonExclusive { ref mut num } => {
                assert!(*num >= 1);
                *num -= 1;
                if *num >= 1 {
                    return;
                }
            },
            GpuAccess::Exclusive { ref mut num } => {
                assert!(*num >= 1);
                *num -= 1;
                if *num >= 1 {
                    return;
                }
            },
        };

        *lock = GpuAccess::None;
    }
}

unsafe impl<T: ?Sized, A> TypedBufferAccess for DeviceLocalBuffer<T, A>
    where T: 'static + Send + Sync
{
    type Content = T;
}
<|endoftext|>