<commit_before>use ast::*;
use env::*;
use parser::*;

pub struct Evaler {
    parser: Parser,
    env:    Env

}

#[allow(dead_code)]
impl Evaler {
    pub fn new() -> Evaler {
        Evaler {
            parser: Parser::new(),
            env:    Env::new()
        }
    }

    pub fn eval(&mut self, code: String) -> ExprAst {
        let ast = self.parser.load(code);
        //self.eval_exp(ast, &mut Env::new())
        self.eval_exp(ast)
    }

    fn eval_exp(&mut self, exp: ExprAst) -> ExprAst {
        if exp.is_self() {
            return exp;
        } else if exp.is_symbol() {
            return self.env.lookup(exp).unwrap();
        } else if exp.is_quote() {
            return exp.cdr().car();
        } else if exp.is_assign() {
            return self.eval_assign(exp);
        } else if exp.is_def() {
            return self.eval_def(exp);
        } else if exp.is_begin() {
            return self.eval_begin(exp);
        } else if exp.is_if() {
            return self.eval_if(exp);
        } else if exp.is_lambda() {
            return self.eval_lambda(exp);
        } else if exp.is_and() {
            return self.eval_and(exp);
        } else if exp.is_or() {
            return self.eval_or(exp);
        } else if exp.is_cond() {
            return self.eval_cond(exp);
        } else if exp.is_let() {
            return self.eval_let(exp);
        } else if exp.is_pair() { //app
            return self.eval_app(exp);
        }
        ExprAst::Symbol(SymbolNode::new("OK"))
    }

    fn eval_assign(&mut self, exp: ExprAst) -> ExprAst {
        let var = exp.c("da");
        let val = self.eval_exp(exp.c("dda"));
        self.env.def_var(var, val);
        ExprAst::Symbol(SymbolNode::new("OK"))
    }

    fn eval_def(&mut self, exp: ExprAst) -> ExprAst {
        let var = exp.def_var();
        let val = exp.def_val();
        let val = self.eval_exp(val);
        self.env.def_var(var, val);
        ExprAst::Symbol(SymbolNode::new("OK"))
    }

    fn eval_if(&mut self, exp: ExprAst) -> ExprAst {
        let pred = exp.c("da");
        let blk_t = exp.c("dda");
        let blk_f = exp.c("ddd");
        let res = self.eval_exp(pred);
        if res.is_true() {
            self.eval_exp(blk_t)
        } else {
            if res.is_empty() {
                ExprAst::Bool(BoolNode::new(false))
            } else {
                self.eval_exp(blk_f.car())
            }
        }
    }

    fn eval_and(&mut self, exp: ExprAst) -> ExprAst {
        let mut elems = exp.cdr();
        if elems.is_empty() {
            return ExprAst::Bool(BoolNode::new(true));
        }
        loop {
            if elems.is_last() { break; }
            let res = self.eval_exp(elems.car());
            if res.is_false() {
                return ExprAst::Bool(BoolNode::new(false));
            }
            elems = elems.cdr();
        }
        self.eval_exp(elems.car())
    }

    fn eval_or(&mut self, exp: ExprAst) -> ExprAst {
        let mut elems = exp.cdr();
        if elems.is_empty() {
            return ExprAst::Bool(BoolNode::new(true));
        }
        loop {
            if elems.is_last() { break; }
            let res = self.eval_exp(elems.car());
            if res.is_true() {
                return ExprAst::Bool(BoolNode::new(true));
            }
            elems = elems.cdr();
        }
        self.eval_exp(elems.car())
    }

    fn eval_cond(&mut self, exp: ExprAst) -> ExprAst {
        let mut elems = exp.cdr().car();
        loop {
            if elems.is_empty() { break; }
            let cur = elems.car();
            let val = self.eval_exp(cur.clone());
            if val.is_true() || val == ExprAst::Symbol(SymbolNode::new("else")) {
                return self.eval_exp(cur.clone().cdr().car());
            }
            elems = elems.cdr();
        }
        ExprAst::Bool(BoolNode::new(true))
    }

    fn eval_let(&mut self, exp: ExprAst) -> ExprAst {
        fn bind_params(exp: ExprAst) -> ExprAst {
            if exp.is_empty() {
                ExprAst::Nil
            } else {
                ExprAst::Pair(PairNode::new(exp.c("aa"), bind_params(exp.cdr())))
            }
        }

        fn bind_argus(exp: ExprAst) -> ExprAst {
            if exp.is_empty() {
                ExprAst::Nil
            } else {
                ExprAst::Pair(PairNode::new(exp.c("ada"), bind_argus(exp.cdr())))
            }
        }

        //FIXME : remove clone?
        let bindings = exp.c("da");
        bindings.print();
        let obj = ExprAst::Pair(PairNode::new(
            bind_params(bindings.clone()).make_lambda(exp.c("dd")),
            bind_argus(bindings)));
        self.eval_exp(obj)
    }

    fn eval_begin(&mut self, exp: ExprAst) -> ExprAst {
        let mut _exp = exp.cdr();
        loop {
            if _exp.is_last() { break; }
            self.eval_exp(_exp.car());
            _exp = _exp.cdr();
        }
        self.eval_exp(_exp.car())
    }

    fn eval_values(&mut self, exprs: ExprAst) -> ExprAst {
        if exprs.is_empty() {
            ExprAst::Nil
        } else {
            let first = self.eval_exp(exprs.car());
            ExprAst::Pair(PairNode::new(first, self.eval_values(exprs.cdr())))
        }
    }

    fn eval_app(&mut self, expr: ExprAst) -> ExprAst {
        let _proc = self.eval_exp(expr.car());
        let _args = self.eval_values(expr.cdr());
        if _proc.is_proc() {
            let func = _proc.as_proc().func();
            return func(_args);
        } else {
            assert!(_proc.is_cproc());
            let _vars = _proc.params();
            self.env = self.env.extend(_vars, _args);
            let begin = ExprAst::Symbol(SymbolNode::new("begin"));
            let res = self.eval_exp(ExprAst::Pair(PairNode::new(begin, _proc.body())));
            self.env = self.env.parent().unwrap();
            res
        }
    }

    fn eval_lambda(&mut self, expr: ExprAst) -> ExprAst {
        //FIXME: remove clone
        ExprAst::CompProc(CompProcNode::new(expr.c("da"),  //vars
                                            expr.c("dd"),  //body
                                            Box::new(self.env.clone())))

    }
}

fn test_evaler() {
    macro_rules! test_case {
        ($test_str:expr, $expect_type:ident, $expect_val:expr) => { {
            let mut evaler = Evaler::new();
            let res = evaler.eval($test_str.to_string());
            if res.$expect_type() != $expect_val {
                assert!(false);
            }
        }}
    }
    test_case!("11", as_int, 11);
    test_case!("'a", as_str, "a");
    test_case!(r#""hello""#, as_str, "hello");
    test_case!("#t", as_bool, true);
    test_case!("(and #t #t)", as_bool, true);
    test_case!("(and #t #f)", as_bool, false);
    test_case!("(or #t #t)", as_bool, true);
    test_case!("(or #f #f)", as_bool, false);
    test_case!("(or #t #f)", as_bool, true);
    test_case!("(set! a 1)", as_str, "OK");
    test_case!("(integer? 1)", as_bool, true);
    test_case!("(boolean? #t)", as_bool, true);
    test_case!("(boolean? #f)", as_bool, true);
    test_case!("(integer? #t)", as_bool, false);
    test_case!("(+ 1 1)", as_int, 2);
    test_case!("(- 1 1)", as_int, 0);
    test_case!("(+ 1 1  1)", as_int, 3);
    test_case!("(+ 1 1 -1 -1)", as_int, 0);
    test_case!("(/ 2 1)", as_int, 2);
    test_case!("(* 2 2)", as_int, 4);
    test_case!("(> 2 1)", as_bool, true);
    test_case!("(> 1 2)", as_bool, false);
    test_case!("(> 1 1)", as_bool, false);
    test_case!("(eq? 1 1)", as_bool, true);
    test_case!("(eq? 1 2)", as_bool, false);
    test_case!("(eq? 1 #f)", as_bool, false);
    test_case!("(pair? '())", as_bool, true);
    test_case!("(> (+ 1 1) 0)", as_bool, true);
    test_case!("(if (> 1 0) 1 else 2)", as_int, 1);
    test_case!("(begin (set! a 1) a)", as_int, 1);
    test_case!("(let ((a 1)) (+ a 1))", as_int, 2);
    test_case!("(let ((a 1)) (> (+ a 1) 0))", as_bool, true);
    test_case!("(let ((a 1) (b 2)) (> a b))", as_bool, false);
    test_case!("(define (add a b) (+ a b))", as_str, "OK");
    test_case!("(cond ((eq? 1 1) 1) ((> 1 2) 2) ))", as_int, 1);
    test_case!("(begin (set! x 5) (set! x 4) (+ x 1))", as_int, 5);
    test_case!("(car '(1 2))", as_int, 1);
    test_case!("(car (cdr '(1 2)))", as_int, 2);
    test_case!("(car (cons 1 2))", as_int, 1);
    test_case!("(car (cdr (cons 1 2)))", as_int, 2);
    test_case!("(car (car (cons (cons 2 3) (cons 1 2))))", as_int, 2);
    test_case!("((lambda (x) x) 1)", as_int, 1);
    test_case!("((lambda (x y) (+ x y )) 1 2)", as_int, 3);
    test_case!("(define add4 (let ((x 4)) (lambda (y) (+ x y))))", as_str, "OK");
    test_case!("((lambda ( x y ) ( if ( = y 0) 1 (* y (x x (- y 1)))))
                (lambda ( x y ) ( if ( = y 0) 1 (* y (x x (- y 1))))) 6)", as_int, 720);

}
<commit_msg>cleanup<commit_after>use ast::*;
use env::*;
use parser::*;

pub struct Evaler {
    parser: Parser,
    env:    Env

}

#[allow(dead_code)]
impl Evaler {
    pub fn new() -> Evaler {
        Evaler {
            parser: Parser::new(),
            env:    Env::new()
        }
    }

    pub fn eval(&mut self, code: String) -> ExprAst {
        let ast = self.parser.load(code);
        self.eval_exp(ast)
    }

    fn eval_exp(&mut self, exp: ExprAst) -> ExprAst {
        if exp.is_self() {
            return exp;
        } else if exp.is_symbol() {
            return self.env.lookup(exp).unwrap();
        } else if exp.is_quote() {
            return exp.cdr().car();
        } else if exp.is_assign() {
            return self.eval_assign(exp);
        } else if exp.is_def() {
            return self.eval_def(exp);
        } else if exp.is_begin() {
            return self.eval_begin(exp);
        } else if exp.is_if() {
            return self.eval_if(exp);
        } else if exp.is_lambda() {
            return self.eval_lambda(exp);
        } else if exp.is_and() {
            return self.eval_and(exp);
        } else if exp.is_or() {
            return self.eval_or(exp);
        } else if exp.is_cond() {
            return self.eval_cond(exp);
        } else if exp.is_let() {
            return self.eval_let(exp);
        } else if exp.is_pair() { //app
            return self.eval_app(exp);
        }
        ExprAst::Symbol(SymbolNode::new("OK"))
    }

    fn eval_assign(&mut self, exp: ExprAst) -> ExprAst {
        let var = exp.c("da");
        let val = self.eval_exp(exp.c("dda"));
        self.env.def_var(var, val);
        ExprAst::Symbol(SymbolNode::new("OK"))
    }

    fn eval_def(&mut self, exp: ExprAst) -> ExprAst {
        let var = exp.def_var();
        let val = exp.def_val();
        let val = self.eval_exp(val);
        self.env.def_var(var, val);
        ExprAst::Symbol(SymbolNode::new("OK"))
    }

    fn eval_if(&mut self, exp: ExprAst) -> ExprAst {
        let pred = exp.c("da");
        let blk_t = exp.c("dda");
        let blk_f = exp.c("ddd");
        let res = self.eval_exp(pred);
        if res.is_true() {
            self.eval_exp(blk_t)
        } else {
            if res.is_empty() {
                ExprAst::Bool(BoolNode::new(false))
            } else {
                self.eval_exp(blk_f.car())
            }
        }
    }

    fn eval_and(&mut self, exp: ExprAst) -> ExprAst {
        let mut elems = exp.cdr();
        if elems.is_empty() {
            return ExprAst::Bool(BoolNode::new(true));
        }
        loop {
            if elems.is_last() { break; }
            let res = self.eval_exp(elems.car());
            if res.is_false() {
                return ExprAst::Bool(BoolNode::new(false));
            }
            elems = elems.cdr();
        }
        self.eval_exp(elems.car())
    }

    fn eval_or(&mut self, exp: ExprAst) -> ExprAst {
        let mut elems = exp.cdr();
        if elems.is_empty() {
            return ExprAst::Bool(BoolNode::new(true));
        }
        loop {
            if elems.is_last() { break; }
            let res = self.eval_exp(elems.car());
            if res.is_true() {
                return ExprAst::Bool(BoolNode::new(true));
            }
            elems = elems.cdr();
        }
        self.eval_exp(elems.car())
    }

    fn eval_cond(&mut self, exp: ExprAst) -> ExprAst {
        let mut elems = exp.cdr().car();
        loop {
            if elems.is_empty() { break; }
            let cur = elems.car();
            let val = self.eval_exp(cur.clone());
            if val.is_true() || val == ExprAst::Symbol(SymbolNode::new("else")) {
                return self.eval_exp(cur.clone().cdr().car());
            }
            elems = elems.cdr();
        }
        ExprAst::Bool(BoolNode::new(true))
    }

    fn eval_let(&mut self, exp: ExprAst) -> ExprAst {
        fn bind_params(exp: ExprAst) -> ExprAst {
            if exp.is_empty() {
                ExprAst::Nil
            } else {
                ExprAst::Pair(PairNode::new(exp.c("aa"), bind_params(exp.cdr())))
            }
        }

        fn bind_args(exp: ExprAst) -> ExprAst {
            if exp.is_empty() {
                ExprAst::Nil
            } else {
                ExprAst::Pair(PairNode::new(exp.c("ada"), bind_args(exp.cdr())))
            }
        }

        //FIXME : remove clone?
        let bindings = exp.c("da");
        bindings.print();
        let obj = ExprAst::Pair(PairNode::new(
            bind_params(bindings.clone()).make_lambda(exp.c("dd")),
            bind_args(bindings)));
        self.eval_exp(obj)
    }

    fn eval_begin(&mut self, exp: ExprAst) -> ExprAst {
        let mut _exp = exp.cdr();
        loop {
            if _exp.is_last() { break; }
            self.eval_exp(_exp.car());
            _exp = _exp.cdr();
        }
        self.eval_exp(_exp.car())
    }

    fn eval_values(&mut self, exprs: ExprAst) -> ExprAst {
        if exprs.is_empty() {
            ExprAst::Nil
        } else {
            let first = self.eval_exp(exprs.car());
            ExprAst::Pair(PairNode::new(first, self.eval_values(exprs.cdr())))
        }
    }

    fn eval_app(&mut self, expr: ExprAst) -> ExprAst {
        let _proc = self.eval_exp(expr.car());
        let _args = self.eval_values(expr.cdr());
        if _proc.is_proc() {
            let func = _proc.as_proc().func();
            return func(_args);
        } else {
            assert!(_proc.is_cproc());
            let _vars = _proc.params();
            self.env = self.env.extend(_vars, _args);
            let begin = ExprAst::Symbol(SymbolNode::new("begin"));
            let res = self.eval_exp(ExprAst::Pair(PairNode::new(begin, _proc.body())));
            self.env = self.env.parent().unwrap();
            res
        }
    }

    fn eval_lambda(&mut self, expr: ExprAst) -> ExprAst {
        //FIXME: remove clone
        ExprAst::CompProc(CompProcNode::new(expr.c("da"),  //vars
                                            expr.c("dd"),  //body
                                            Box::new(self.env.clone())))

    }
}

fn test_evaler() {
    macro_rules! test_case {
        ($test_str:expr, $expect_type:ident, $expect_val:expr) => { {
            let mut evaler = Evaler::new();
            let res = evaler.eval($test_str.to_string());
            if res.$expect_type() != $expect_val {
                assert!(false);
            }
        }}
    }
    test_case!("11", as_int, 11);
    test_case!("'a", as_str, "a");
    test_case!(r#""hello""#, as_str, "hello");
    test_case!("#t", as_bool, true);
    test_case!("(and #t #t)", as_bool, true);
    test_case!("(and #t #f)", as_bool, false);
    test_case!("(or #t #t)", as_bool, true);
    test_case!("(or #f #f)", as_bool, false);
    test_case!("(or #t #f)", as_bool, true);
    test_case!("(set! a 1)", as_str, "OK");
    test_case!("(integer? 1)", as_bool, true);
    test_case!("(boolean? #t)", as_bool, true);
    test_case!("(boolean? #f)", as_bool, true);
    test_case!("(integer? #t)", as_bool, false);
    test_case!("(+ 1 1)", as_int, 2);
    test_case!("(- 1 1)", as_int, 0);
    test_case!("(+ 1 1  1)", as_int, 3);
    test_case!("(+ 1 1 -1 -1)", as_int, 0);
    test_case!("(/ 2 1)", as_int, 2);
    test_case!("(* 2 2)", as_int, 4);
    test_case!("(> 2 1)", as_bool, true);
    test_case!("(> 1 2)", as_bool, false);
    test_case!("(> 1 1)", as_bool, false);
    test_case!("(eq? 1 1)", as_bool, true);
    test_case!("(eq? 1 2)", as_bool, false);
    test_case!("(eq? 1 #f)", as_bool, false);
    test_case!("(pair? '())", as_bool, true);
    test_case!("(> (+ 1 1) 0)", as_bool, true);
    test_case!("(if (> 1 0) 1 else 2)", as_int, 1);
    test_case!("(begin (set! a 1) a)", as_int, 1);
    test_case!("(let ((a 1)) (+ a 1))", as_int, 2);
    test_case!("(let ((a 1)) (> (+ a 1) 0))", as_bool, true);
    test_case!("(let ((a 1) (b 2)) (> a b))", as_bool, false);
    test_case!("(define (add a b) (+ a b))", as_str, "OK");
    test_case!("(cond ((eq? 1 1) 1) ((> 1 2) 2) ))", as_int, 1);
    test_case!("(begin (set! x 5) (set! x 4) (+ x 1))", as_int, 5);
    test_case!("(car '(1 2))", as_int, 1);
    test_case!("(car (cdr '(1 2)))", as_int, 2);
    test_case!("(car (cons 1 2))", as_int, 1);
    test_case!("(car (cdr (cons 1 2)))", as_int, 2);
    test_case!("(car (car (cons (cons 2 3) (cons 1 2))))", as_int, 2);
    test_case!("((lambda (x) x) 1)", as_int, 1);
    test_case!("((lambda (x y) (+ x y )) 1 2)", as_int, 3);
    test_case!("(define add4 (let ((x 4)) (lambda (y) (+ x y))))", as_str, "OK");
    test_case!("((lambda ( x y ) ( if ( = y 0) 1 (* y (x x (- y 1)))))
                (lambda ( x y ) ( if ( = y 0) 1 (* y (x x (- y 1))))) 6)", as_int, 720);

}
<|endoftext|>


<commit_before>use std::fmt::{Display, Error, Formatter};
use std::rand;

pub struct Game {
    intended_word: String,
    guessed: Vec<GuessedRes>,
}

impl Game {
    pub fn new(words: Vec<String>) -> Game {
        let intended_word = words.get(rand::random() % words.len()).unwrap().clone();
        let guessed = intended_word.chars().map(|ch| GuessedRes::new(ch, false)).collect();
        Game {
            intended_word: intended_word,
            guessed : guessed,
        }
    }

    pub fn intended_word<'r>(&'r self) -> &'r str {
        &self.intended_word[]
    }

    pub fn current_progress(&self) -> CurrentProgress {
        CurrentProgress { vec : self.guessed.iter().map(|x| x.to_enum()).collect() }
    }

    pub fn letters_left(&self) -> usize {
        self.guessed.iter().filter(|&x| !x.is_guessed).count()
    }

    pub fn guess(&mut self, guess : &str) {
        for ch in guess.chars() {
            self.guess_char(ch);
        }
    }

    fn guess_char(&mut self, ch : char) {
        for g_res in self.guessed.iter_mut().filter(|x| x.ch==ch) {
            g_res.is_guessed = true;
        }
    }
}

struct GuessedRes {
    ch : char,
    is_guessed : bool,
}

impl GuessedRes {

    fn new(ch : char, is_guessed : bool) -> GuessedRes {
        GuessedRes {
            ch : ch,
            is_guessed : is_guessed,
        }
    }

    fn to_enum(&self) -> Guessed {
        if self.is_guessed {
            Some(self.ch)
        } else {
            None
        }
    }
}

pub type Guessed = Option<char>;

#[derive(Debug)]
pub struct CurrentProgress {
    vec : Vec<Guessed>,
}

impl Display for CurrentProgress {
    fn fmt(&self, f : &mut Formatter) -> Result<(), Error> {
        let mut result = String::with_capacity(self.vec.len() * 2);
        for guess in self.vec.iter() {
            match *guess {
                Some(ch) => result.push(ch),
                None => result.push('_'),
            }
            result.push(' ');
        }
        write!(f, "{}", result.as_slice())
    }
}
<commit_msg>Переименование GuessedRes в CharState<commit_after>use std::fmt::{Display, Error, Formatter};
use std::rand;

pub struct Game {
    intended_word: String,
    guessed: Vec<CharState>,
}

impl Game {
    pub fn new(words: Vec<String>) -> Game {
        let intended_word = words.get(rand::random() % words.len()).unwrap().clone();
        let guessed = intended_word.chars().map(|ch| CharState::new(ch, false)).collect();
        Game {
            intended_word: intended_word,
            guessed : guessed,
        }
    }

    pub fn intended_word<'r>(&'r self) -> &'r str {
        &self.intended_word[]
    }

    pub fn current_progress(&self) -> CurrentProgress {
        CurrentProgress { vec : self.guessed.iter().map(|x| x.to_enum()).collect() }
    }

    pub fn letters_left(&self) -> usize {
        self.guessed.iter().filter(|&x| !x.is_guessed).count()
    }

    pub fn guess(&mut self, guess : &str) {
        for ch in guess.chars() {
            self.guess_char(ch);
        }
    }

    fn guess_char(&mut self, ch : char) {
        for g_res in self.guessed.iter_mut().filter(|x| x.ch==ch) {
            g_res.is_guessed = true;
        }
    }
}

struct CharState {
    ch : char,
    is_guessed : bool,
}

impl CharState {

    fn new(ch : char, is_guessed : bool) -> CharState {
        CharState {
            ch : ch,
            is_guessed : is_guessed,
        }
    }

    fn to_enum(&self) -> Guessed {
        if self.is_guessed {
            Some(self.ch)
        } else {
            None
        }
    }
}

pub type Guessed = Option<char>;

#[derive(Debug)]
pub struct CurrentProgress {
    vec : Vec<Guessed>,
}

impl Display for CurrentProgress {
    fn fmt(&self, f : &mut Formatter) -> Result<(), Error> {
        let mut result = String::with_capacity(self.vec.len() * 2);
        for guess in self.vec.iter() {
            match *guess {
                Some(ch) => result.push(ch),
                None => result.push('_'),
            }
            result.push(' ');
        }
        write!(f, "{}", result.as_slice())
    }
}
<|endoftext|>


<commit_before>trait JoinIterator {
  fn join(&mut self, sep: &str) -> String;
}

impl<A: ToString, T: Iterator<A>> JoinIterator for T {
  fn join(&mut self, sep: &str) -> String {
    match self.next() {
      Some(elem) => {
        let mut output = String::from_str(elem.to_string().as_slice());
        for elem in *self {
          output.push_str(sep);
          output.push_str(elem.to_string().as_slice());
        }
        output
      }
      None => String::new()
    }
  }
}

#[test]
fn test_join() {
  let many = vec![1u,2,3];
  let one  = vec![1u];
  let none: Vec<uint> = vec![];

  assert_eq!(many.iter().join(", ").as_slice(), "1, 2, 3");
  assert_eq!(one .iter().join(", ").as_slice(), "1");
  assert_eq!(none.iter().join(", ").as_slice(), "");
}
<commit_msg>add comments<commit_after>// An Iterator adapter that walks through all the elements in the Iterator,
// converts them to Strings and joins them to one big String, seperated by
// some seperator string slice.
trait JoinIterator {
  fn join(&mut self, sep: &str) -> String;
}

// Implement it for all Iterators with Elements convertable into a String
impl<A: ToString, T: Iterator<A>> JoinIterator for T {
  fn join(&mut self, sep: &str) -> String {
    match self.next() {
      Some(elem) => {
        let mut output = String::from_str(elem.to_string().as_slice());
        for elem in *self {
          output.push_str(sep);
          output.push_str(elem.to_string().as_slice());
        }
        output
      }
      None => String::new()
    }
  }
}

#[test]
fn test_join() {
  let many = vec![1u,2,3];
  let one  = vec![1u];
  let none: Vec<uint> = vec![];

  assert_eq!(many.iter().join(", ").as_slice(), "1, 2, 3");
  assert_eq!(one .iter().join(", ").as_slice(), "1");
  assert_eq!(none.iter().join(", ").as_slice(), "");
}
<|endoftext|>


<commit_before>use endian_type::{BigEndian, LittleEndian};
use std::ffi::OsString;
use std::path::{Path, PathBuf};

use nibble_vec::Nibblet;

/// Trait for types which can be used to key a Radix Trie.
///
/// Types that implement this trait should be convertible to a vector of half-bytes (nibbles)
/// such that no two instances of the type convert to the same vector.
/// To protect against faulty behaviour, the trie will **panic** if it finds two distinct keys
/// of type `K` which encode to the same `Nibblet`, so be careful!
///
/// If you would like to implement this trait for your own type, you need to implement
/// *either* `encode_bytes` or `encode`. You only need to implement one of the two.
/// If you don't implement one, your code will **panic** as soon you use the trie.
/// There is no performance penalty for implementing `encode_bytes` instead of `encode`,
/// so it is preferred except in the case where you require half-byte precision.
///
/// Many standard types implement this trait already. Integer types are encoded *big-endian*
/// by default but can be encoded little-endian using the `LittleEndian<T>` wrapper type.
pub trait TrieKey: PartialEq + Eq {
    /// Encode a value as a vector of bytes.
    fn encode_bytes(&self) -> Vec<u8> {
        panic!("implement this method or TrieKey::encode");
    }

    /// Encode a value as a NibbleVec.
    #[inline]
    fn encode(&self) -> Nibblet {
        Nibblet::from_byte_vec(self.encode_bytes())
    }
}

/// Key comparison result.
#[derive(Debug)]
pub enum KeyMatch {
    /// The keys match up to the given index.
    Partial(usize),
    /// The first key is a prefix of the second.
    FirstPrefix,
    /// The second key is a prefix of the first.
    SecondPrefix,
    /// The keys match exactly.
    Full,
}

/// Compare two Trie keys.
///
/// Compares `first[start_idx .. ]` to `second`, i.e. only looks at a slice of the first key.
#[inline]
pub fn match_keys(start_idx: usize, first: &Nibblet, second: &Nibblet) -> KeyMatch {
    let first_len = first.len() - start_idx;
    let min_length = ::std::cmp::min(first_len, second.len());

    for i in 0..min_length {
        if first.get(start_idx + i) != second.get(i) {
            return KeyMatch::Partial(i);
        }
    }

    match (first_len, second.len()) {
        (x, y) if x < y => KeyMatch::FirstPrefix,
        (x, y) if x == y => KeyMatch::Full,
        _ => KeyMatch::SecondPrefix,
    }
}

/// Check two keys for equality and panic if they differ.
#[inline]
pub fn check_keys<K: ?Sized>(key1: &K, key2: &K)
where
    K: TrieKey,
{
    if *key1 != *key2 {
        panic!("multiple-keys with the same bit representation.");
    }
}

// --- TrieKey Implementations for standard types --- ///

// This blanket implementation goes into play when specialization is stabilized
// impl<T> TrieKey for T where T: Into<Vec<u8>> + Clone + Eq + PartialEq {
// fn encode_bytes(&self) -> Vec<u8> {
// self.clone().into()
// }
// }

impl TrieKey for Vec<u8> {
    #[inline]
    fn encode_bytes(&self) -> Vec<u8> {
        self.clone()
    }
}

impl TrieKey for [u8] {
    #[inline]
    fn encode_bytes(&self) -> Vec<u8> {
        self.to_vec()
    }
}

impl TrieKey for String {
    #[inline]
    fn encode_bytes(&self) -> Vec<u8> {
        self.as_bytes().encode_bytes()
    }
}

impl TrieKey for str {
    #[inline]
    fn encode_bytes(&self) -> Vec<u8> {
        self.as_bytes().encode_bytes()
    }
}

impl<'a, T: ?Sized + TrieKey> TrieKey for &'a T {
    #[inline]
    fn encode_bytes(&self) -> Vec<u8> {
        (**self).encode_bytes()
    }
}

impl<'a, T: ?Sized + TrieKey> TrieKey for &'a mut T {
    #[inline]
    fn encode_bytes(&self) -> Vec<u8> {
        (**self).encode_bytes()
    }
}

impl TrieKey for i8 {
    #[inline]
    fn encode_bytes(&self) -> Vec<u8> {
        let mut v: Vec<u8> = Vec::with_capacity(1);
        v.push(*self as u8);
        v
    }
}

impl TrieKey for u8 {
    #[inline]
    fn encode_bytes(&self) -> Vec<u8> {
        let mut v: Vec<u8> = Vec::with_capacity(1);
        v.push(*self);
        v
    }
}

#[cfg(unix)]
impl TrieKey for PathBuf {
    fn encode_bytes(&self) -> Vec<u8> {
        use std::os::unix::ffi::OsStringExt;
        let str: OsString = self.clone().into();
        str.into_vec()
    }
}

#[cfg(unix)]
impl TrieKey for Path {
    fn encode_bytes(&self) -> Vec<u8> {
        use std::os::unix::ffi::OsStrExt;
        self.as_os_str().as_bytes().encode_bytes()
    }
}

impl<T> TrieKey for LittleEndian<T>
where
    T: Eq + Copy,
{
    fn encode_bytes(&self) -> Vec<u8> {
        self.as_bytes().encode_bytes()
    }
}

impl<T> TrieKey for BigEndian<T>
where
    T: Eq + Copy,
{
    fn encode_bytes(&self) -> Vec<u8> {
        self.as_bytes().to_vec()
    }
}

macro_rules! int_keys {
    ( $( $t:ty ),* ) => {
        $(
        impl TrieKey for $t {
            fn encode_bytes(&self) -> Vec<u8> {
                let be: BigEndian<$t> = From::from(*self);
                be.encode_bytes()
            }
        }
        )*
    };
}

int_keys!(u16, u32, u64, i16, i32, i64, usize, isize);

#[cfg(test)]
mod test {
    pub trait DefaultTrieKey {
        fn encode_bytes(&self) -> Vec<u8>;
    }

    impl<T: Into<Vec<u8>> + Clone + PartialEq + Eq> DefaultTrieKey for T {
        #[inline]
        fn encode_bytes(&self) -> Vec<u8> {
            self.clone().into()
        }
    }

    pub trait AsTrieKey {
        fn encode_bytes(&self) -> Vec<u8>;
    }

    impl<T: AsRef<[u8]> + Clone + PartialEq + Eq> AsTrieKey for &T {
        #[inline]
        fn encode_bytes(&self) -> Vec<u8> {
            self.as_ref().to_vec()
        }
    }

    macro_rules! encode_bytes {
        ($e:expr) => {
            (&$e).encode_bytes()
        };
    }

    #[test]
    fn test_autoref_specialization() {
        let _ = encode_bytes!([0_u8]);
        let _ = encode_bytes!("hello");
        let _ = encode_bytes!("hello".to_string());
    }
}
<commit_msg>TrieKey support for Vec<T> of integer types (#63)<commit_after>use endian_type::{BigEndian, LittleEndian};
use std::ffi::OsString;
use std::path::{Path, PathBuf};

use nibble_vec::Nibblet;

/// Trait for types which can be used to key a Radix Trie.
///
/// Types that implement this trait should be convertible to a vector of half-bytes (nibbles)
/// such that no two instances of the type convert to the same vector.
/// To protect against faulty behaviour, the trie will **panic** if it finds two distinct keys
/// of type `K` which encode to the same `Nibblet`, so be careful!
///
/// If you would like to implement this trait for your own type, you need to implement
/// *either* `encode_bytes` or `encode`. You only need to implement one of the two.
/// If you don't implement one, your code will **panic** as soon you use the trie.
/// There is no performance penalty for implementing `encode_bytes` instead of `encode`,
/// so it is preferred except in the case where you require half-byte precision.
///
/// Many standard types implement this trait already. Integer types are encoded *big-endian*
/// by default but can be encoded little-endian using the `LittleEndian<T>` wrapper type.
pub trait TrieKey: PartialEq + Eq {
    /// Encode a value as a vector of bytes.
    fn encode_bytes(&self) -> Vec<u8> {
        panic!("implement this method or TrieKey::encode");
    }

    /// Encode a value as a NibbleVec.
    #[inline]
    fn encode(&self) -> Nibblet {
        Nibblet::from_byte_vec(self.encode_bytes())
    }
}

/// Key comparison result.
#[derive(Debug)]
pub enum KeyMatch {
    /// The keys match up to the given index.
    Partial(usize),
    /// The first key is a prefix of the second.
    FirstPrefix,
    /// The second key is a prefix of the first.
    SecondPrefix,
    /// The keys match exactly.
    Full,
}

/// Compare two Trie keys.
///
/// Compares `first[start_idx .. ]` to `second`, i.e. only looks at a slice of the first key.
#[inline]
pub fn match_keys(start_idx: usize, first: &Nibblet, second: &Nibblet) -> KeyMatch {
    let first_len = first.len() - start_idx;
    let min_length = ::std::cmp::min(first_len, second.len());

    for i in 0..min_length {
        if first.get(start_idx + i) != second.get(i) {
            return KeyMatch::Partial(i);
        }
    }

    match (first_len, second.len()) {
        (x, y) if x < y => KeyMatch::FirstPrefix,
        (x, y) if x == y => KeyMatch::Full,
        _ => KeyMatch::SecondPrefix,
    }
}

/// Check two keys for equality and panic if they differ.
#[inline]
pub fn check_keys<K: ?Sized>(key1: &K, key2: &K)
where
    K: TrieKey,
{
    if *key1 != *key2 {
        panic!("multiple-keys with the same bit representation.");
    }
}

// --- TrieKey Implementations for standard types --- ///

// This blanket implementation goes into play when specialization is stabilized
// impl<T> TrieKey for T where T: Into<Vec<u8>> + Clone + Eq + PartialEq {
// fn encode_bytes(&self) -> Vec<u8> {
// self.clone().into()
// }
// }

impl TrieKey for Vec<u8> {
    #[inline]
    fn encode_bytes(&self) -> Vec<u8> {
        self.clone()
    }
}

impl TrieKey for [u8] {
    #[inline]
    fn encode_bytes(&self) -> Vec<u8> {
        self.to_vec()
    }
}

impl TrieKey for String {
    #[inline]
    fn encode_bytes(&self) -> Vec<u8> {
        self.as_bytes().encode_bytes()
    }
}

impl TrieKey for str {
    #[inline]
    fn encode_bytes(&self) -> Vec<u8> {
        self.as_bytes().encode_bytes()
    }
}

impl<'a, T: ?Sized + TrieKey> TrieKey for &'a T {
    #[inline]
    fn encode_bytes(&self) -> Vec<u8> {
        (**self).encode_bytes()
    }
}

impl<'a, T: ?Sized + TrieKey> TrieKey for &'a mut T {
    #[inline]
    fn encode_bytes(&self) -> Vec<u8> {
        (**self).encode_bytes()
    }
}

impl TrieKey for i8 {
    #[inline]
    fn encode_bytes(&self) -> Vec<u8> {
        let mut v: Vec<u8> = Vec::with_capacity(1);
        v.push(*self as u8);
        v
    }
}

impl TrieKey for u8 {
    #[inline]
    fn encode_bytes(&self) -> Vec<u8> {
        let mut v: Vec<u8> = Vec::with_capacity(1);
        v.push(*self);
        v
    }
}

#[cfg(unix)]
impl TrieKey for PathBuf {
    fn encode_bytes(&self) -> Vec<u8> {
        use std::os::unix::ffi::OsStringExt;
        let str: OsString = self.clone().into();
        str.into_vec()
    }
}

#[cfg(unix)]
impl TrieKey for Path {
    fn encode_bytes(&self) -> Vec<u8> {
        use std::os::unix::ffi::OsStrExt;
        self.as_os_str().as_bytes().encode_bytes()
    }
}

impl<T> TrieKey for LittleEndian<T>
where
    T: Eq + Copy,
{
    fn encode_bytes(&self) -> Vec<u8> {
        self.as_bytes().encode_bytes()
    }
}

impl<T> TrieKey for BigEndian<T>
where
    T: Eq + Copy,
{
    fn encode_bytes(&self) -> Vec<u8> {
        self.as_bytes().to_vec()
    }
}

macro_rules! int_keys {
    ( $( $t:ty ),* ) => {
        $(
        impl TrieKey for $t {
            fn encode_bytes(&self) -> Vec<u8> {
                let be: BigEndian<$t> = From::from(*self);
                be.encode_bytes()
            }
        }
        )*
    };
}

int_keys!(u16, u32, u64, i16, i32, i64, usize, isize);

macro_rules! vec_int_keys {
  ( $( $t:ty ),* ) => {
      $(
      impl TrieKey for Vec<$t> {
          fn encode_bytes(&self) -> Vec<u8> {
              let mut v = Vec::<u8>::with_capacity(self.len() * std::mem::size_of::<$t>());
              for u in self {
                  v.extend_from_slice(&u.to_be_bytes());
              }
              v
          }
      }
      )*
   };
}

vec_int_keys!(u16, u32, u64, i16, i32, i64, usize, isize);

#[cfg(test)]
mod test {
    pub trait DefaultTrieKey {
        fn encode_bytes(&self) -> Vec<u8>;
    }

    impl<T: Into<Vec<u8>> + Clone + PartialEq + Eq> DefaultTrieKey for T {
        #[inline]
        fn encode_bytes(&self) -> Vec<u8> {
            self.clone().into()
        }
    }

    pub trait AsTrieKey {
        fn encode_bytes(&self) -> Vec<u8>;
    }

    impl<T: AsRef<[u8]> + Clone + PartialEq + Eq> AsTrieKey for &T {
        #[inline]
        fn encode_bytes(&self) -> Vec<u8> {
            self.as_ref().to_vec()
        }
    }

    macro_rules! encode_bytes {
        ($e:expr) => {
            (&$e).encode_bytes()
        };
    }

    #[test]
    fn test_autoref_specialization() {
        let _ = encode_bytes!([0_u8]);
        let _ = encode_bytes!("hello");
        let _ = encode_bytes!("hello".to_string());
    }
}
<|endoftext|>


<commit_before>use position::Position;
use cppn::{Cppn, CppnNodeType};
use acyclic_network::NodeType;
use std::fmt::Debug;

/// Represents a node in the substrate. `T` stores additional information about that node.
#[derive(Clone, Debug)]
pub struct Node<P, T>
    where P: Position,
          T: NodeType
{
    pub position: P,
    pub node_type: T,
}

#[derive(Clone, Debug)]
pub struct Layer<P, T>
    where P: Position,
          T: NodeType
{
    nodes: Vec<Node<P, T>>,
}

impl<P, T> Layer<P, T>
    where P: Position,
          T: NodeType
{
    pub fn new() -> Self {
        Layer { nodes: Vec::new() }
    }

    pub fn nodes(&self) -> &[Node<P, T>] {
        &self.nodes
    }

    pub fn add_node(&mut self, position: P, node_type: T) {
        self.nodes.push(Node {
            position: position,
            node_type: node_type,
        });
    }
}

#[derive(Clone)]
pub struct Link<'a, P, T>
    where P: Position + 'a,
          T: NodeType + 'a
{
    pub source: &'a Node<P, T>,
    pub target: &'a Node<P, T>,
    pub source_idx: usize,
    pub target_idx: usize,
    pub outputs: Vec<f64>,
    pub distance: f64,
}

#[derive(Clone, Debug)]
pub enum LinkMode {
    AbsolutePositions,
    AbsolutePositionsAndDistance,
    RelativePositionOfTarget,
}

#[derive(Clone, Debug)]
struct LayerLink {
    from_layer: usize,
    to_layer: usize,
    max_distance: Option<f64>,
}

#[derive(Clone, Debug)]
pub struct Substrate<P, T>
    where P: Position,
          T: NodeType
{
    layers: Vec<Layer<P, T>>,
    layer_links: Vec<LayerLink>,
}

impl<P, T> Substrate<P, T>
    where P: Position,
          T: NodeType
{
    pub fn new() -> Self {
        Substrate {
            layers: Vec::new(),
            layer_links: Vec::new(),
        }
    }

    pub fn layers(&self) -> &[Layer<P, T>] {
        &self.layers
    }

    pub fn add_layer(&mut self, layer: Layer<P, T>) -> usize {
        let layer_idx = self.layers.len();
        self.layers.push(layer);
        return layer_idx;
    }

    pub fn add_layer_link(&mut self,
                          from_layer: usize,
                          to_layer: usize,
                          max_distance: Option<f64>) {
        self.layer_links.push(LayerLink {
            from_layer: from_layer,
            to_layer: to_layer,
            max_distance: max_distance,
        });
    }

    /// Iterate over all produced links of Cppn.

    pub fn each_link<'a, N, L, EXTID, F>(&'a self,
                                         cppn: &'a mut Cppn<'a, N, L, EXTID>,
                                         mode: LinkMode,
                                         callback: &mut F)
        where N: CppnNodeType,
              L: Copy + Debug + Send + Sized + Into<f64> + 'a,
              EXTID: Copy + Debug + Send + Sized + Ord + 'a,
              F: FnMut(Link<'a, P, T>)
    {
        for layer_link in self.layer_links.iter() {
            for (source_idx, source) in self.layers[layer_link.from_layer]
                                            .nodes
                                            .iter()
                                            .enumerate() {
                // Reject invalid connections.
                if !source.node_type.accept_outgoing_links() {
                    continue;
                }

                for (target_idx, target) in self.layers[layer_link.to_layer]
                                                .nodes
                                                .iter()
                                                .enumerate() {
                    // Reject invalid connections.
                    if !target.node_type.accept_incoming_links() {
                        continue;
                    }

                    let distance = source.position.distance(&target.position);

                    // reject a pair of nodes based on `max_distance`.
                    if let Some(max_d) = layer_link.max_distance {
                        if distance > max_d {
                            continue;
                        }
                    }

                    // Calculate the weight between source and target using the CPPN.

                    let outputs_from_cppn = match mode {
                        LinkMode::AbsolutePositions => {
                            let inputs_to_cppn = [source.position.coords(),
                                                  target.position.coords()];
                            cppn.calculate(&inputs_to_cppn)
                        }
                        LinkMode::AbsolutePositionsAndDistance => {
                            let distance_between: &[_] = &[distance];
                            let inputs_to_cppn = [source.position.coords(),
                                                  target.position.coords(),
                                                  distance_between];
                            cppn.calculate(&inputs_to_cppn)
                        }
                        LinkMode::RelativePositionOfTarget => {
                            let relative_position_of_target =
                                source.position.relative_position(&target.position);
                            let inputs_to_cppn = [source.position.coords(),
                                                  relative_position_of_target.coords()];
                            cppn.calculate(&inputs_to_cppn)
                        }
                    };

                    let link = Link {
                        source: source,
                        target: target,
                        source_idx: source_idx,
                        target_idx: target_idx,
                        outputs: outputs_from_cppn,
                        distance: distance,
                    };
                    callback(link);
                }
            }
        }
    }
}
<commit_msg>Include the layer index<commit_after>use position::Position;
use cppn::{Cppn, CppnNodeType};
use acyclic_network::NodeType;
use std::fmt::Debug;

/// Represents a node in the substrate. `T` stores additional information about that node.
#[derive(Clone, Debug)]
pub struct Node<P, T>
    where P: Position,
          T: NodeType
{
    pub position: P,
    pub node_type: T,
}

#[derive(Clone, Debug)]
pub struct Layer<P, T>
    where P: Position,
          T: NodeType
{
    nodes: Vec<Node<P, T>>,
}

impl<P, T> Layer<P, T>
    where P: Position,
          T: NodeType
{
    pub fn new() -> Self {
        Layer { nodes: Vec::new() }
    }

    pub fn nodes(&self) -> &[Node<P, T>] {
        &self.nodes
    }

    pub fn add_node(&mut self, position: P, node_type: T) {
        self.nodes.push(Node {
            position: position,
            node_type: node_type,
        });
    }
}

#[derive(Clone)]
pub struct Link<'a, P, T>
    where P: Position + 'a,
          T: NodeType + 'a
{
    pub source: &'a Node<P, T>,
    pub target: &'a Node<P, T>,
    pub source_idx: (usize, usize), // (layer, node)
    pub target_idx: (usize, usize), // (layer, node)
    pub outputs: Vec<f64>,
    pub distance: f64,
}

#[derive(Clone, Debug)]
pub enum LinkMode {
    AbsolutePositions,
    AbsolutePositionsAndDistance,
    RelativePositionOfTarget,
}

#[derive(Clone, Debug)]
struct LayerLink {
    from_layer: usize,
    to_layer: usize,
    max_distance: Option<f64>,
}

#[derive(Clone, Debug)]
pub struct Substrate<P, T>
    where P: Position,
          T: NodeType
{
    layers: Vec<Layer<P, T>>,
    layer_links: Vec<LayerLink>,
}

impl<P, T> Substrate<P, T>
    where P: Position,
          T: NodeType
{
    pub fn new() -> Self {
        Substrate {
            layers: Vec::new(),
            layer_links: Vec::new(),
        }
    }

    pub fn layers(&self) -> &[Layer<P, T>] {
        &self.layers
    }

    pub fn add_layer(&mut self, layer: Layer<P, T>) -> usize {
        let layer_idx = self.layers.len();
        self.layers.push(layer);
        return layer_idx;
    }

    pub fn add_layer_link(&mut self,
                          from_layer: usize,
                          to_layer: usize,
                          max_distance: Option<f64>) {
        self.layer_links.push(LayerLink {
            from_layer: from_layer,
            to_layer: to_layer,
            max_distance: max_distance,
        });
    }

    /// Iterate over all produced links of Cppn.

    pub fn each_link<'a, N, L, EXTID, F>(&'a self,
                                         cppn: &'a mut Cppn<'a, N, L, EXTID>,
                                         mode: LinkMode,
                                         callback: &mut F)
        where N: CppnNodeType,
              L: Copy + Debug + Send + Sized + Into<f64> + 'a,
              EXTID: Copy + Debug + Send + Sized + Ord + 'a,
              F: FnMut(Link<'a, P, T>)
    {
        for layer_link in self.layer_links.iter() {
            for (source_idx, source) in self.layers[layer_link.from_layer]
                                            .nodes
                                            .iter()
                                            .enumerate() {
                // Reject invalid connections.
                if !source.node_type.accept_outgoing_links() {
                    continue;
                }

                for (target_idx, target) in self.layers[layer_link.to_layer]
                                                .nodes
                                                .iter()
                                                .enumerate() {
                    // Reject invalid connections.
                    if !target.node_type.accept_incoming_links() {
                        continue;
                    }

                    let distance = source.position.distance(&target.position);

                    // reject a pair of nodes based on `max_distance`.
                    if let Some(max_d) = layer_link.max_distance {
                        if distance > max_d {
                            continue;
                        }
                    }

                    // Calculate the weight between source and target using the CPPN.

                    let outputs_from_cppn = match mode {
                        LinkMode::AbsolutePositions => {
                            let inputs_to_cppn = [source.position.coords(),
                                                  target.position.coords()];
                            cppn.calculate(&inputs_to_cppn)
                        }
                        LinkMode::AbsolutePositionsAndDistance => {
                            let distance_between: &[_] = &[distance];
                            let inputs_to_cppn = [source.position.coords(),
                                                  target.position.coords(),
                                                  distance_between];
                            cppn.calculate(&inputs_to_cppn)
                        }
                        LinkMode::RelativePositionOfTarget => {
                            let relative_position_of_target =
                                source.position.relative_position(&target.position);
                            let inputs_to_cppn = [source.position.coords(),
                                                  relative_position_of_target.coords()];
                            cppn.calculate(&inputs_to_cppn)
                        }
                    };

                    let link = Link {
                        source: source,
                        target: target,
                        source_idx: (layer_link.from_layer, source_idx),
                        target_idx: (layer_link.to_layer, target_idx),
                        outputs: outputs_from_cppn,
                        distance: distance,
                    };
                    callback(link);
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::fs;
use std::io::{self, Write};
use std::path::Path;
use std::process::{Command, Stdio};
use walkdir::WalkDir;

use util::*;

actor!{
    #[derive(Debug)]
    pub struct Tarballer {
        /// The input folder to be compressed
        input: String = "package",

        /// The prefix of the tarballs
        output: String = "./dist",

        /// The fold in which the input is to be found
        work_dir: String = "./workdir",
    }
}

impl Tarballer {
    /// Generate the actual tarballs
    pub fn run(self) -> io::Result<()> {
        let path = get_path()?;
        need_cmd(&path, "tar")?;
        need_cmd(&path, "gzip")?;
        let have_xz = need_either_cmd(&path, "xz", "7z")?;

        let tar = self.output + ".tar";
        let tar_gz = tar.clone() + ".gz";
        let tar_xz = tar.clone() + ".xz";

        // Remove any existing files
        for file in &[&tar, &tar_gz, &tar_xz] {
            if Path::new(file).exists() {
                fs::remove_file(file)?;
            }
        }

        // Sort files by their suffix, to group files with the same name from
        // different locations (likely identical) and files with the same
        // extension (likely containing similar data).
        let mut paths = get_recursive_paths(self.work_dir.as_ref(), self.input.as_ref())?;
        paths.sort_by(|a, b| a.bytes().rev().cmp(b.bytes().rev()));

        // Write the tar file
        let mut child = Command::new("tar")
            .arg("-cf")
            .arg(&tar)
            .arg("-T")
            .arg("-")
            .stdin(Stdio::piped())
            .current_dir(&self.work_dir)
            .spawn()?;
        if let Some(stdin) = child.stdin.as_mut() {
            for path in paths {
                writeln!(stdin, "{}", path)?;
            }
        }
        let status = child.wait()?;
        if !status.success() {
            let msg = format!("failed to make tarball: {}", status);
            return Err(io::Error::new(io::ErrorKind::Other, msg));
        }

        // Write the .tar.xz file
        let status = if have_xz {
            Command::new("xz")
                .arg("-9")
                .arg("--keep")
                .arg(&tar)
                .status()?
        } else {
            Command::new("7z")
                .arg("a")
                .arg("-bd")
                .arg("-txz")
                .arg("-mx=9")
                .arg("-mmt=off")
                .arg(&tar_xz)
                .arg(&tar)
                .status()?
        };
        if !status.success() {
            let msg = format!("failed to make tar.xz: {}", status);
            return Err(io::Error::new(io::ErrorKind::Other, msg));
        }

        // Write the .tar.gz file (removing the .tar)
        let status = Command::new("gzip")
            .arg(&tar)
            .status()?;
        if !status.success() {
            let msg = format!("failed to make tar.gz: {}", status);
            return Err(io::Error::new(io::ErrorKind::Other, msg));
        }

        Ok(())
    }
}

fn get_recursive_paths(root: &Path, name: &Path) -> io::Result<Vec<String>> {
    let mut paths = vec![];
    for entry in WalkDir::new(root.join(name)).min_depth(1) {
        let entry = entry?;
        let path = entry.path().strip_prefix(root).unwrap();
        let path = path.to_str().unwrap().to_owned();

        if entry.file_type().is_dir() {
            // Include only empty dirs, as others get add via their contents.
            // FIXME: do we really need empty dirs at all?
            if fs::read_dir(entry.path())?.next().is_none() {
                paths.push(path);
            }
        } else {
            paths.push(path);
        }
    }
    Ok(paths)
}
<commit_msg>Use the tar crate for creating tarballs<commit_after>// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::fs;
use std::io;
use std::path::Path;
use std::process::Command;

use tar::Builder;
use walkdir::WalkDir;

use util::*;

actor!{
    #[derive(Debug)]
    pub struct Tarballer {
        /// The input folder to be compressed
        input: String = "package",

        /// The prefix of the tarballs
        output: String = "./dist",

        /// The fold in which the input is to be found
        work_dir: String = "./workdir",
    }
}

impl Tarballer {
    /// Generate the actual tarballs
    pub fn run(self) -> io::Result<()> {
        let path = get_path()?;
        need_cmd(&path, "gzip")?;
        let have_xz = need_either_cmd(&path, "xz", "7z")?;

        let tar = self.output + ".tar";
        let tar_gz = tar.clone() + ".gz";
        let tar_xz = tar.clone() + ".xz";

        // Remove any existing files
        for file in &[&tar, &tar_gz, &tar_xz] {
            if Path::new(file).exists() {
                fs::remove_file(file)?;
            }
        }

        // Sort files by their suffix, to group files with the same name from
        // different locations (likely identical) and files with the same
        // extension (likely containing similar data).
        let mut paths = get_recursive_paths(self.work_dir.as_ref(), self.input.as_ref())?;
        paths.sort_by(|a, b| a.bytes().rev().cmp(b.bytes().rev()));

        // Write the tar file
        let output = fs::File::create(&tar)?;
        let mut builder = Builder::new(output);
        for path in paths {
            let path = Path::new(&path);
            let src = Path::new(&self.work_dir).join(path);
            if path.is_dir() {
                builder.append_dir(path, src)?;
            } else {
                let mut src = fs::File::open(src)?;
                builder.append_file(path, &mut src)?;
            }
        }
        builder.into_inner()?;

        // Write the .tar.xz file
        let status = if have_xz {
            Command::new("xz")
                .arg("-9")
                .arg("--keep")
                .arg(&tar)
                .status()?
        } else {
            Command::new("7z")
                .arg("a")
                .arg("-bd")
                .arg("-txz")
                .arg("-mx=9")
                .arg("-mmt=off")
                .arg(&tar_xz)
                .arg(&tar)
                .status()?
        };
        if !status.success() {
            let msg = format!("failed to make tar.xz: {}", status);
            return Err(io::Error::new(io::ErrorKind::Other, msg));
        }

        // Write the .tar.gz file (removing the .tar)
        let status = Command::new("gzip")
            .arg(&tar)
            .status()?;
        if !status.success() {
            let msg = format!("failed to make tar.gz: {}", status);
            return Err(io::Error::new(io::ErrorKind::Other, msg));
        }

        Ok(())
    }
}

fn get_recursive_paths(root: &Path, name: &Path) -> io::Result<Vec<String>> {
    let mut paths = vec![];
    for entry in WalkDir::new(root.join(name)).min_depth(1) {
        let entry = entry?;
        let path = entry.path().strip_prefix(root).unwrap();
        let path = path.to_str().unwrap().to_owned();

        if entry.file_type().is_dir() {
            // Include only empty dirs, as others get add via their contents.
            // FIXME: do we really need empty dirs at all?
            if fs::read_dir(entry.path())?.next().is_none() {
                paths.push(path);
            }
        } else {
            paths.push(path);
        }
    }
    Ok(paths)
}
<|endoftext|>


<commit_before>extern crate procrs;
extern crate argparse;
#[macro_use]
extern crate prettytable;
use prettytable::Table;
use prettytable::format::FormatBuilder;
use std::collections::HashMap;
use std::iter::repeat;
use procrs::pid::*;
use procrs::TaskId;
use argparse::{ArgumentParser, StoreTrue, Store};

fn main() {
    let opts = parse_args();
    let (query, long, perf, verbose, tree) =
        (opts.query, opts.long, opts.perf, opts.verbose, opts.tree);

    let mut pids = PidIter::new_query(query).unwrap()
        .collect::<Result<Vec<_>, _>>().unwrap();

    let mut name_indent = HashMap::new();

    if verbose {
        for pid in pids {
            println!("{:?}", pid);
        }
        return
    }

    if opts.tree {
        pids = treeify_names(pids, &mut name_indent);
    } else {
        pids.sort_by(|p1, p2| p1.stat.pid.cmp(&p2.stat.pid));
    };
    // Assume hertz is 100.
    // TODO: Look this up via syscall (no /proc value for it)
    let hertz = 100;
    let minute_hertz = hertz * 60;
    let hour_hertz = minute_hertz * 60;

    let mut table = Table::init(
        pids.iter().map(|p| {
            // When we have a tree, the name is prepended with an indent.
            let mut name = match tree {
                false => String::new(),
                true => name_indent.remove(&p.stat.pid).unwrap()
            };

            // For long output, try using the cmdline first.
            // FIXME: Sometimes prog_name != cmdline[0].
            if !long {
                name.push_str(&p.stat.comm);
            } else {
                let cmdline = p.cmdline.join(" ");
                name.push_str(
                    match cmdline {
                         ref s if s.len() > 0 => s,
                        _ => &p.stat.comm
                    }
                );
            }

            match (long, perf) {
                (_, false) => {
                    row![p.stat.pid, p.stat.ppid, name]
                },
                (_, true) => {
                    let rss = p.status.vmrss.map(|m| (m / 1024).to_string()).unwrap_or("".to_owned());
                    // FIXME: This algorithm is definitely wrong
                    let second_utime = p.stat.utime % hertz;
                    let minute_utime = p.stat.utime / minute_hertz;
                    let hour_utime = p.stat.utime / hour_hertz;
                    let utime = format!(
                        "{:02}:{:02}:{:02}",
                        hour_utime,
                        minute_utime,
                        second_utime
                    );
                    row![p.stat.pid, p.stat.ppid, rss, utime, name]
                }
            }
        }).collect::<Vec<_>>()
    );

    table.set_titles(
        match (long, perf) {
            (false, false) =>
                row!["Pid", "Ppid", "Cmd"],
            (true, false) =>
                row!["Pid", "Ppid", "Cmd"],
            (false, true) =>
                row!["Pid", "Ppid", "RSS", "Time", "Cmd"],
            (true, true) =>
                row!["Pid", "Ppid", "RSS", "Time", "Cmd"]
        }
    );
    table.set_format(
        FormatBuilder::new()
            .column_separator(' ')
            .build()
    );
    table.printstd();
}

// Given a vector of Pid structs, treeify their names, and return them in the right order.
// This is similar to ps -AH.
fn treeify_names(pids: Vec<Pid>, name_indents: &mut HashMap<TaskId, String>) -> Vec<Pid> {
    let mut child_pids = HashMap::new();
    for pid in pids {
        let ppid = pid.stat.ppid;
        child_pids.entry(ppid)
            .or_insert(Vec::new())
            .push(pid);
    }
    enumerate_children(0, &mut child_pids, name_indents, -1)
}

// Enumerate children pids, and return them.
fn enumerate_children(pid: TaskId, child_pids: &mut HashMap<TaskId, Vec<Pid>>,
    name_indents: &mut HashMap<TaskId, String>, indent: i32) -> Vec<Pid> {
    name_indents.insert(pid,
        match indent {
            i if i >= 0 =>
                repeat("  ").take(i as usize).collect::<String>(),
            _ => "".to_owned()
        }
    );
    let mut pids = Vec::new();
    let ppids = match child_pids.remove(&pid) {
        Some(v) => v,
        None => { return pids; }
    };
    for pid in ppids {
        let pid_num = pid.stat.pid;
        pids.push(pid);
        pids.append(
            &mut enumerate_children(pid_num, child_pids, name_indents, indent + 1)
        );
    }
    pids
}

struct ProgOpts {
    query: PidQuery,
    tree: bool,
    perf: bool,
    long: bool,
    verbose: bool
}

fn parse_args() -> ProgOpts {
    let mut opts = ProgOpts {
        query: PidQuery::NoneQuery,
        tree: false,
        perf: false,
        long: false,
        verbose: false
    };

    {
        let mut ap = ArgumentParser::new();
        ap.set_description("Query linux processes");
        ap.refer(&mut opts.tree)
            .add_option(&["-t", "--tree"], StoreTrue, "Display process tree");
        ap.refer(&mut opts.perf)
            .add_option(&["-p", "--perf"], StoreTrue, "Display performance information");
        ap.refer(&mut opts.long)
            .add_option(&["-l", "--long"], StoreTrue, "Display more information");
        ap.refer(&mut opts.verbose)
            .add_option(&["-v", "--verbose"], StoreTrue, "Verbose output");
        ap.refer(&mut opts.query)
            .add_argument("query", Store, "Optional query to search by, pid or string");
        ap.parse_args_or_exit();
    }

    opts
}
<commit_msg>Attempt at fixing cputime algorithm.<commit_after>extern crate procrs;
extern crate argparse;
#[macro_use]
extern crate prettytable;
use prettytable::Table;
use prettytable::format::FormatBuilder;
use std::collections::HashMap;
use std::iter::repeat;
use procrs::pid::*;
use procrs::TaskId;
use argparse::{ArgumentParser, StoreTrue, Store};

fn main() {
    let opts = parse_args();
    let (query, long, perf, verbose, tree) =
        (opts.query, opts.long, opts.perf, opts.verbose, opts.tree);

    let mut pids = PidIter::new_query(query).unwrap()
        .collect::<Result<Vec<_>, _>>().unwrap();

    let mut name_indent = HashMap::new();

    if verbose {
        for pid in pids {
            println!("{:?}", pid);
        }
        return
    }

    if opts.tree {
        pids = treeify_names(pids, &mut name_indent);
    } else {
        pids.sort_by(|p1, p2| p1.stat.pid.cmp(&p2.stat.pid));
    };
    // Assume hertz is 100.
    // TODO: Look this up via syscall (no /proc value for it)
    let hertz = 100;
    // let minute_hertz = hertz * 60;
    // let hour_hertz = minute_hertz * 60;

    let mut table = Table::init(
        pids.iter().map(|p| {
            // When we have a tree, the name is prepended with an indent.
            let mut name = match tree {
                false => String::new(),
                true => name_indent.remove(&p.stat.pid).unwrap()
            };

            // For long output, try using the cmdline first.
            // FIXME: Sometimes prog_name != cmdline[0].
            if !long {
                name.push_str(&p.stat.comm);
            } else {
                let cmdline = p.cmdline.join(" ");
                name.push_str(
                    match cmdline {
                         ref s if s.len() > 0 => s,
                        _ => &p.stat.comm
                    }
                );
            }

            match (long, perf) {
                (_, false) => {
                    row![p.stat.pid, p.stat.ppid, name]
                },
                (_, true) => {
                    let rss = p.status.vmrss.map(|m| (m / 1024).to_string()).unwrap_or("".to_owned());
                    // FIXME: This algorithm is definitely wrong
                    let _second_utime = p.stat.utime / hertz;
                    let second_utime = _second_utime % 60;
                    let _minute_utime = _second_utime / 60;
                    let minute_utime = _minute_utime % 60;
                    let hour_utime = _minute_utime / 60;
                    let utime = format!(
                        "{:02}:{:02}:{:02}",
                        hour_utime,
                        minute_utime,
                        second_utime
                    );
                    row![p.stat.pid, p.stat.ppid, rss, utime, name]
                }
            }
        }).collect::<Vec<_>>()
    );

    table.set_titles(
        match (long, perf) {
            (false, false) =>
                row!["Pid", "Ppid", "Cmd"],
            (true, false) =>
                row!["Pid", "Ppid", "Cmd"],
            (false, true) =>
                row!["Pid", "Ppid", "RSS", "Time", "Cmd"],
            (true, true) =>
                row!["Pid", "Ppid", "RSS", "Time", "Cmd"]
        }
    );
    table.set_format(
        FormatBuilder::new()
            .column_separator(' ')
            .build()
    );
    table.printstd();
}

// Given a vector of Pid structs, treeify their names, and return them in the right order.
// This is similar to ps -AH.
fn treeify_names(pids: Vec<Pid>, name_indents: &mut HashMap<TaskId, String>) -> Vec<Pid> {
    let mut child_pids = HashMap::new();
    for pid in pids {
        let ppid = pid.stat.ppid;
        child_pids.entry(ppid)
            .or_insert(Vec::new())
            .push(pid);
    }
    enumerate_children(0, &mut child_pids, name_indents, -1)
}

// Enumerate children pids, and return them.
fn enumerate_children(pid: TaskId, child_pids: &mut HashMap<TaskId, Vec<Pid>>,
    name_indents: &mut HashMap<TaskId, String>, indent: i32) -> Vec<Pid> {
    name_indents.insert(pid,
        match indent {
            i if i >= 0 =>
                repeat("  ").take(i as usize).collect::<String>(),
            _ => "".to_owned()
        }
    );
    let mut pids = Vec::new();
    let ppids = match child_pids.remove(&pid) {
        Some(v) => v,
        None => { return pids; }
    };
    for pid in ppids {
        let pid_num = pid.stat.pid;
        pids.push(pid);
        pids.append(
            &mut enumerate_children(pid_num, child_pids, name_indents, indent + 1)
        );
    }
    pids
}

struct ProgOpts {
    query: PidQuery,
    tree: bool,
    perf: bool,
    long: bool,
    verbose: bool
}

fn parse_args() -> ProgOpts {
    let mut opts = ProgOpts {
        query: PidQuery::NoneQuery,
        tree: false,
        perf: false,
        long: false,
        verbose: false
    };

    {
        let mut ap = ArgumentParser::new();
        ap.set_description("Query linux processes");
        ap.refer(&mut opts.tree)
            .add_option(&["-t", "--tree"], StoreTrue, "Display process tree");
        ap.refer(&mut opts.perf)
            .add_option(&["-p", "--perf"], StoreTrue, "Display performance information");
        ap.refer(&mut opts.long)
            .add_option(&["-l", "--long"], StoreTrue, "Display more information");
        ap.refer(&mut opts.verbose)
            .add_option(&["-v", "--verbose"], StoreTrue, "Verbose output");
        ap.refer(&mut opts.query)
            .add_argument("query", Store, "Optional query to search by, pid or string");
        ap.parse_args_or_exit();
    }

    opts
}
<|endoftext|>


<commit_before>use super::ring_key::RingKey;
use crate::{crypto::{hash,
                     keys::{Permissioned,
                            ToKeyString},
                     SECRET_SYM_KEY_SUFFIX},
            error::{Error,
                    Result},
            fs::AtomicWriter};
use std::{convert::TryInto,
          io::Write,
          path::{Path,
                 PathBuf}};

pub struct KeyCache(PathBuf);

impl KeyCache {
    pub fn new<P>(path: P) -> Self
        where P: Into<PathBuf>
    {
        KeyCache(path.into())
    }

    /// Ensure that the directory backing the cache exists on disk.
    pub fn setup(&self) -> Result<()> {
        if !self.0.is_dir() {
            std::fs::create_dir_all(&self.0)?;
        }
        Ok(())
    }

    pub fn write_ring_key(&self, key: &RingKey) -> Result<()> { self.maybe_write_key(key) }

    /// Note: name is just the name, not the name + revision
    pub fn latest_ring_key_revision(&self, name: &str) -> Result<RingKey> {
        match self.get_latest_path_for(name, SECRET_SYM_KEY_SUFFIX)? {
            Some(path) => path.try_into(),
            None => {
                let msg = format!("No revisions found for {} ring key", name);
                Err(Error::CryptoError(msg))
            }
        }
    }

    /// Provides the path at which this file would be found in the
    /// cache, if it exists (or, alternatively, where it would be
    /// written to).
    fn path_in_cache<K>(&self, key: K) -> PathBuf
        where K: AsRef<Path>
    {
        self.0.join(key.as_ref())
    }

    /// Write the given key into the cache. If the key already exists,
    /// and the content has the same hash value, nothing will be
    /// done. If the file exists and it has *different* content, an
    /// Error is returned.
    fn maybe_write_key<K>(&self, key: &K) -> Result<()>
        where K: AsRef<Path> + ToKeyString + Permissioned
    {
        let keyfile = self.path_in_cache(&key);
        let content = key.to_key_string()?;

        if keyfile.is_file() {
            let existing_hash = hash::hash_file(&keyfile)?;
            let new_hash = hash::hash_string(&content);
            if existing_hash != new_hash {
                let msg = format!("Existing key file {} found but new version hash is different, \
                                   failing to write new file over existing. (existing = {}, \
                                   incoming = {})",
                                  keyfile.display(),
                                  existing_hash,
                                  new_hash);
                return Err(Error::CryptoError(msg));
            }
        } else {
            // Technically speaking, this probably doesn't really need
            // to be an atomic write process, since we just tested
            // that the file doesn't currently exist. It does,
            // however, bundle up writing with platform-independent
            // permission setting, which is *super* convenient.
            let w = AtomicWriter::new_with_permissions(&keyfile, K::permissions())?;
            w.with_writer(|f| f.write_all(content.as_ref()))?;
        }
        Ok(())
    }

    /// Search the key cache for all files that are revisions of the
    /// given key. Returns the full paths to those files within the
    /// cache.
    fn get_all_paths_for(&self,
                         name: &str,
                         key_extension: &str)
                         -> Result<impl Iterator<Item = PathBuf>> {
        // Ideally, we'd want that `*` to be `\d{14}` to match the
        // structure of our revisions... perhaps that can be an
        // additional filter later on with an actual regex?
        let pattern = self.0.join(format!("{}-*.{}", name, key_extension));
        let pattern = pattern.to_string_lossy();

        // TODO (CM): this is a bogus error
        Ok(glob::glob(&pattern).map_err(|_e| Error::CryptoError("Couldn't glob!".to_string()))?
                               .filter_map(std::result::Result::ok)
                               .filter(|p| p.metadata().map(|m| m.is_file()).unwrap_or(false)))
    }

    /// Given a key name and extension, find the path that corresponds
    /// to the most recent revision of that key in the cache, if it
    /// exists.
    fn get_latest_path_for(&self, name: &str, key_extension: &str) -> Result<Option<PathBuf>> {
        Ok(self.get_all_paths_for(name, key_extension)?.max())
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::crypto::test_support::*;

    static VALID_KEY: &str = "ring-key-valid-20160504220722.sym.key";
    static VALID_NAME_WITH_REV: &str = "ring-key-valid-20160504220722";

    #[test]
    fn get_all_paths_for() {
        let (cache, _dir) = new_cache();

        let paths = cache.get_all_paths_for("beyonce", SECRET_SYM_KEY_SUFFIX)
                         .unwrap();
        assert_eq!(paths.count(), 0);

        cache.write_ring_key(&RingKey::new("beyonce")).unwrap();
        let paths = cache.get_all_paths_for("beyonce", SECRET_SYM_KEY_SUFFIX)
                         .unwrap();
        assert_eq!(paths.count(), 1);

        wait_1_sec(); // ensure new revision
                      // will be different.
        cache.write_ring_key(&RingKey::new("beyonce")).unwrap();

        let paths = cache.get_all_paths_for("beyonce", SECRET_SYM_KEY_SUFFIX)
                         .unwrap();
        assert_eq!(paths.count(), 2);

        // We should not include another named key in the count
        cache.write_ring_key(&RingKey::new("jayz")).unwrap();
        let paths = cache.get_all_paths_for("beyonce", SECRET_SYM_KEY_SUFFIX)
                         .unwrap();
        assert_eq!(paths.count(), 2);
    }

    #[test]
    fn latest_cached_revision_single() {
        let (cache, _dir) = new_cache();

        let key = RingKey::new("beyonce");
        cache.write_ring_key(&key).unwrap();

        let latest = cache.latest_ring_key_revision("beyonce").unwrap();
        assert_eq!(latest.name(), key.name());
        assert_eq!(latest.revision(), key.revision());
    }

    #[test]
    fn latest_cached_revision_multiple() {
        let (cache, _dir) = new_cache();

        let k1 = RingKey::new("beyonce");
        cache.write_ring_key(&k1).unwrap();

        wait_1_sec();

        let k2 = RingKey::new("beyonce");
        cache.write_ring_key(&k2).unwrap();

        assert_eq!(k1.name(), k2.name());
        assert_ne!(k1.revision(), k2.revision());

        let latest = cache.latest_ring_key_revision("beyonce").unwrap();
        assert_eq!(latest.name(), k2.name());
        assert_eq!(latest.revision(), k2.revision());
    }

    #[test]
    #[should_panic(expected = "No revisions found for")]
    fn latest_cached_revision_nonexistent() {
        let (cache, _dir) = new_cache();
        cache.latest_ring_key_revision("nope-nope").unwrap();
    }

    #[test]
    fn writing_ring_key() {
        let (cache, dir) = new_cache();

        let content = fixture_as_string(&format!("keys/{}", VALID_KEY));
        let new_key_file = dir.path().join(VALID_KEY);
        assert_eq!(new_key_file.is_file(), false);

        let key: RingKey = content.parse().unwrap();
        assert_eq!(key.name_with_rev(), VALID_NAME_WITH_REV);
        cache.write_ring_key(&key).unwrap();
        assert!(new_key_file.is_file());

        let new_content = std::fs::read_to_string(new_key_file).unwrap();
        assert_eq!(new_content, content);
    }

    #[test]
    fn write_key_with_existing_identical() {
        let (cache, dir) = new_cache();
        let content = fixture_as_string(&format!("keys/{}", VALID_KEY));
        let new_key_file = dir.path().join(VALID_KEY);

        // install the key into the cache
        std::fs::copy(fixture(&format!("keys/{}", VALID_KEY)), &new_key_file).unwrap();

        let key: RingKey = content.parse().unwrap();
        cache.write_ring_key(&key).unwrap();
        assert_eq!(key.name_with_rev(), VALID_NAME_WITH_REV);
        assert!(new_key_file.is_file());
    }

    #[test]
    #[should_panic(expected = "Existing key file")]
    fn write_key_exists_but_hashes_differ() {
        let (cache, dir) = new_cache();
        let old_content = fixture_as_string("keys/ring-key-valid-20160504220722.sym.key");

        std::fs::write(dir.path().join("ring-key-valid-20160504220722.sym.key"),
                       &old_content).unwrap();

        #[rustfmt::skip]
        let new_content = "SYM-SEC-1\nring-key-valid-20160504220722\n\nkA+c03Ly5qEoOZIjJ5zCD2vHI05pAW59PfCOb8thmZw=";

        assert_ne!(old_content, new_content);

        let new_key: RingKey = new_content.parse().unwrap();
        // this should fail
        cache.write_ring_key(&new_key).unwrap();
    }

    // Old tests... not fully converting over to new implementation
    // yet because I think the function won't be sticking around very
    // long.

    // #[test]
    // fn cached_path() {
    //     let (cache, dir) = new_cache();
    //     fs::copy(fixture(&format!("keys/{}", VALID_KEY)),
    //              dir.path().join(VALID_KEY)).unwrap();

    //     let result = cache.ring_key_cached_path(VALID_NAME_WITH_REV).unwrap();
    //     assert_eq!(result, cache.path().join(VALID_KEY));
    // }

    // #[test]
    // #[should_panic(expected = "No secret key found at")]
    // fn get_secret_key_path_nonexistent() {
    //     let (cache, _dir) = new_cache();
    //     cache.ring_key_cached_path(VALID_NAME_WITH_REV).unwrap();
    // }
}
<commit_msg>[REFACTOR] Extract fetch_latest_revision generic function for cache<commit_after>use super::ring_key::RingKey;
use crate::{crypto::{hash,
                     keys::{KeyExtension,
                            Permissioned,
                            ToKeyString}},
            error::{Error,
                    Result},
            fs::AtomicWriter};
use std::{convert::TryFrom,
          io::Write,
          path::{Path,
                 PathBuf}};

pub struct KeyCache(PathBuf);

impl KeyCache {
    pub fn new<P>(path: P) -> Self
        where P: Into<PathBuf>
    {
        KeyCache(path.into())
    }

    /// Ensure that the directory backing the cache exists on disk.
    pub fn setup(&self) -> Result<()> {
        if !self.0.is_dir() {
            std::fs::create_dir_all(&self.0)?;
        }
        Ok(())
    }

    pub fn write_ring_key(&self, key: &RingKey) -> Result<()> { self.maybe_write_key(key) }

    /// Note: name is just the name, not the name + revision
    pub fn latest_ring_key_revision(&self, name: &str) -> Result<RingKey> {
        self.fetch_latest_revision::<RingKey>(name)
    }

    ////////////////////////////////////////////////////////////////////////

    // TODO (CM): Turn this into Option<Result<K>>
    fn fetch_latest_revision<K>(&self, name: &str) -> Result<K>
        where K: KeyExtension + TryFrom<PathBuf, Error = Error>
    {
        match self.get_latest_path_for(name, K::extension())? {
            Some(path) => <K as TryFrom<PathBuf>>::try_from(path),
            None => {
                let msg = format!("No revisions found for {}", name);
                Err(Error::CryptoError(msg))
            }
        }
    }

    /// Provides the path at which this file would be found in the
    /// cache, if it exists (or, alternatively, where it would be
    /// written to).
    fn path_in_cache<K>(&self, key: K) -> PathBuf
        where K: AsRef<Path>
    {
        self.0.join(key.as_ref())
    }

    /// Write the given key into the cache. If the key already exists,
    /// and the content has the same hash value, nothing will be
    /// done. If the file exists and it has *different* content, an
    /// Error is returned.
    fn maybe_write_key<K>(&self, key: &K) -> Result<()>
        where K: AsRef<Path> + ToKeyString + Permissioned
    {
        let keyfile = self.path_in_cache(&key);
        let content = key.to_key_string()?;

        if keyfile.is_file() {
            let existing_hash = hash::hash_file(&keyfile)?;
            let new_hash = hash::hash_string(&content);
            if existing_hash != new_hash {
                let msg = format!("Existing key file {} found but new version hash is different, \
                                   failing to write new file over existing. (existing = {}, \
                                   incoming = {})",
                                  keyfile.display(),
                                  existing_hash,
                                  new_hash);
                return Err(Error::CryptoError(msg));
            }
        } else {
            // Technically speaking, this probably doesn't really need
            // to be an atomic write process, since we just tested
            // that the file doesn't currently exist. It does,
            // however, bundle up writing with platform-independent
            // permission setting, which is *super* convenient.
            let w = AtomicWriter::new_with_permissions(&keyfile, K::permissions())?;
            w.with_writer(|f| f.write_all(content.as_ref()))?;
        }
        Ok(())
    }

    /// Search the key cache for all files that are revisions of the
    /// given key. Returns the full paths to those files within the
    /// cache.
    fn get_all_paths_for(&self,
                         name: &str,
                         key_extension: &str)
                         -> Result<impl Iterator<Item = PathBuf>> {
        // Ideally, we'd want that `*` to be `\d{14}` to match the
        // structure of our revisions... perhaps that can be an
        // additional filter later on with an actual regex?
        let pattern = self.0.join(format!("{}-*.{}", name, key_extension));
        let pattern = pattern.to_string_lossy();

        // TODO (CM): this is a bogus error
        Ok(glob::glob(&pattern).map_err(|_e| Error::CryptoError("Couldn't glob!".to_string()))?
                               .filter_map(std::result::Result::ok)
                               .filter(|p| p.metadata().map(|m| m.is_file()).unwrap_or(false)))
    }

    /// Given a key name and extension, find the path that corresponds
    /// to the most recent revision of that key in the cache, if it
    /// exists.
    fn get_latest_path_for(&self, name: &str, key_extension: &str) -> Result<Option<PathBuf>> {
        Ok(self.get_all_paths_for(name, key_extension)?.max())
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::crypto::test_support::*;

    static VALID_KEY: &str = "ring-key-valid-20160504220722.sym.key";
    static VALID_NAME_WITH_REV: &str = "ring-key-valid-20160504220722";

    #[test]
    fn get_all_paths_for() {
        let (cache, _dir) = new_cache();

        let paths = cache.get_all_paths_for("beyonce", RingKey::extension())
                         .unwrap();
        assert_eq!(paths.count(), 0);

        cache.write_ring_key(&RingKey::new("beyonce")).unwrap();
        let paths = cache.get_all_paths_for("beyonce", RingKey::extension())
                         .unwrap();
        assert_eq!(paths.count(), 1);

        wait_1_sec(); // ensure new revision
                      // will be different.
        cache.write_ring_key(&RingKey::new("beyonce")).unwrap();

        let paths = cache.get_all_paths_for("beyonce", RingKey::extension())
                         .unwrap();
        assert_eq!(paths.count(), 2);

        // We should not include another named key in the count
        cache.write_ring_key(&RingKey::new("jayz")).unwrap();
        let paths = cache.get_all_paths_for("beyonce", RingKey::extension())
                         .unwrap();
        assert_eq!(paths.count(), 2);
    }

    #[test]
    fn latest_cached_revision_single() {
        let (cache, _dir) = new_cache();

        let key = RingKey::new("beyonce");
        cache.write_ring_key(&key).unwrap();

        let latest = cache.latest_ring_key_revision("beyonce").unwrap();
        assert_eq!(latest.name(), key.name());
        assert_eq!(latest.revision(), key.revision());
    }

    #[test]
    fn latest_cached_revision_multiple() {
        let (cache, _dir) = new_cache();

        let k1 = RingKey::new("beyonce");
        cache.write_ring_key(&k1).unwrap();

        wait_1_sec();

        let k2 = RingKey::new("beyonce");
        cache.write_ring_key(&k2).unwrap();

        assert_eq!(k1.name(), k2.name());
        assert_ne!(k1.revision(), k2.revision());

        let latest = cache.latest_ring_key_revision("beyonce").unwrap();
        assert_eq!(latest.name(), k2.name());
        assert_eq!(latest.revision(), k2.revision());
    }

    #[test]
    #[should_panic(expected = "No revisions found for")]
    fn latest_cached_revision_nonexistent() {
        let (cache, _dir) = new_cache();
        cache.latest_ring_key_revision("nope-nope").unwrap();
    }

    #[test]
    fn writing_ring_key() {
        let (cache, dir) = new_cache();

        let content = fixture_as_string(&format!("keys/{}", VALID_KEY));
        let new_key_file = dir.path().join(VALID_KEY);
        assert_eq!(new_key_file.is_file(), false);

        let key: RingKey = content.parse().unwrap();
        assert_eq!(key.name_with_rev(), VALID_NAME_WITH_REV);
        cache.write_ring_key(&key).unwrap();
        assert!(new_key_file.is_file());

        let new_content = std::fs::read_to_string(new_key_file).unwrap();
        assert_eq!(new_content, content);
    }

    #[test]
    fn write_key_with_existing_identical() {
        let (cache, dir) = new_cache();
        let content = fixture_as_string(&format!("keys/{}", VALID_KEY));
        let new_key_file = dir.path().join(VALID_KEY);

        // install the key into the cache
        std::fs::copy(fixture(&format!("keys/{}", VALID_KEY)), &new_key_file).unwrap();

        let key: RingKey = content.parse().unwrap();
        cache.write_ring_key(&key).unwrap();
        assert_eq!(key.name_with_rev(), VALID_NAME_WITH_REV);
        assert!(new_key_file.is_file());
    }

    #[test]
    #[should_panic(expected = "Existing key file")]
    fn write_key_exists_but_hashes_differ() {
        let (cache, dir) = new_cache();
        let old_content = fixture_as_string("keys/ring-key-valid-20160504220722.sym.key");

        std::fs::write(dir.path().join("ring-key-valid-20160504220722.sym.key"),
                       &old_content).unwrap();

        #[rustfmt::skip]
        let new_content = "SYM-SEC-1\nring-key-valid-20160504220722\n\nkA+c03Ly5qEoOZIjJ5zCD2vHI05pAW59PfCOb8thmZw=";

        assert_ne!(old_content, new_content);

        let new_key: RingKey = new_content.parse().unwrap();
        // this should fail
        cache.write_ring_key(&new_key).unwrap();
    }

    // Old tests... not fully converting over to new implementation
    // yet because I think the function won't be sticking around very
    // long.

    // #[test]
    // fn cached_path() {
    //     let (cache, dir) = new_cache();
    //     fs::copy(fixture(&format!("keys/{}", VALID_KEY)),
    //              dir.path().join(VALID_KEY)).unwrap();

    //     let result = cache.ring_key_cached_path(VALID_NAME_WITH_REV).unwrap();
    //     assert_eq!(result, cache.path().join(VALID_KEY));
    // }

    // #[test]
    // #[should_panic(expected = "No secret key found at")]
    // fn get_secret_key_path_nonexistent() {
    //     let (cache, _dir) = new_cache();
    //     cache.ring_key_cached_path(VALID_NAME_WITH_REV).unwrap();
    // }
}
<|endoftext|>