<commit_before>use driver::ctxt::Context;
use error::msg::Msg;

use parser::ast::*;
use parser::ast::Expr::*;
use parser::ast::Stmt::*;
use parser::ast::Type::*;
use parser::ast::visit::Visitor;

use sym::*;
use sym::Sym::*;

pub fn check(ctxt: &Context, ast: &Ast) {
    DefCheck::new(ctxt).visit_ast(ast);

    if ctxt.diag.borrow().has_errors() { return; }

    TypeCheck::new(ctxt).visit_ast(ast);
}

struct DefCheck<'a, 'ast: 'a> {
    ctxt: &'a Context<'a, 'ast>,
    current_id: Option<NodeId>,
}

impl<'a, 'ast> DefCheck<'a, 'ast> {
    fn new(ctxt: &'a Context<'a, 'ast>) -> DefCheck<'a, 'ast> {
        DefCheck {
            ctxt: ctxt,
            current_id: None
        }
    }
}

impl<'a, 'ast> Visitor<'ast> for DefCheck<'a, 'ast> {
    fn visit_fct(&mut self, f: &'ast Function) {
        self.current_id = Some(f.id);

        visit::walk_fct(self, f);

        // if no type for function found, then function returns ()
        self.ctxt.types.borrow_mut().entry(f.id).or_insert(BuiltinType::Unit);
    }

    fn visit_param(&mut self, p: &'ast Param) {
        self.current_id = Some(p.id);

        visit::walk_param(self, p);
    }

    fn visit_stmt(&mut self, s: &'ast Stmt) {
        if let StmtVar(ref var) = *s {
            self.current_id = Some(var.id)
        }

        visit::walk_stmt(self, s);
    }

    fn visit_type(&mut self, t: &'ast Type) {
        match *t {
            TypeBasic(ref basic) => {
                if let Some(builtin) = self.ctxt.sym.borrow().get_type(basic.name) {
                    let id = self.current_id.unwrap();
                    self.ctxt.types.borrow_mut().insert(id, builtin);
                } else {
                    let tyname = self.ctxt.interner.str(basic.name).to_string();
                    let msg = Msg::UnknownType(tyname);
                    self.ctxt.diag.borrow_mut().report(basic.pos, msg);
                }
            }

            _ => self.ctxt.diag.borrow_mut().report_unimplemented(t.pos())
        }
    }
}

struct TypeCheck<'a, 'ast: 'a> {
    ctxt: &'a Context<'a, 'ast>,
    expr_type: BuiltinType,
    fct: Option<&'ast Function>,
}

impl<'a, 'ast> TypeCheck<'a, 'ast> {
    fn new(ctxt: &'a Context<'a, 'ast>) -> TypeCheck<'a, 'ast> {
        TypeCheck {
            ctxt: ctxt,
            expr_type: BuiltinType::Unit,
            fct: None
        }
    }

    fn check_stmt_var(&mut self, s: &'ast StmtVarType) {
        let expr_type = s.expr.as_ref().map(|expr| {
            self.visit_expr(&expr);

            self.expr_type
        });

        let defined_type = if let Some(ref ty) = s.data_type {
            self.ctxt.types.borrow().get(&s.id).map(|bt| *bt)
        } else {
            expr_type
        };

        if defined_type.is_none() {
            let tyname = self.ctxt.interner.str(s.name).to_string();
            self.ctxt.diag.borrow_mut().report(s.pos, Msg::VarNeedsTypeInfo(tyname));

            return;
        }

        if expr_type.is_some() && (defined_type.unwrap() != expr_type.unwrap()) {
            let varname = self.ctxt.interner.str(s.name).to_string();
            let defname = defined_type.unwrap().to_string();
            let exprname = expr_type.unwrap().to_string();
            let msg = Msg::VarTypesIncompatible(varname, defname, exprname);

            self.ctxt.diag.borrow_mut().report(s.pos, msg);
        }
    }

    fn check_stmt_while(&mut self, s: &'ast StmtWhileType) {
        self.visit_expr(&s.cond);

        if self.expr_type != BuiltinType::Bool {
            let tyname = self.expr_type.to_string();
            let msg = Msg::WhileCondType(tyname);

            self.ctxt.diag.borrow_mut().report(s.pos, msg);
        }

        self.visit_stmt(&s.block);
    }

    fn check_stmt_if(&mut self, s: &'ast StmtIfType) {
        self.visit_expr(&s.cond);

        if self.expr_type != BuiltinType::Bool {
            let tyname = self.expr_type.to_string();
            let msg = Msg::IfCondType(tyname);

            self.ctxt.diag.borrow_mut().report(s.pos, msg);
        }

        self.visit_stmt(&s.then_block);

        if let Some(ref else_block) = s.else_block {
            self.visit_stmt(else_block);
        }
    }

    fn check_stmt_return(&mut self, s: &'ast StmtReturnType) {
        let expr_type = s.expr.as_ref().map(|expr| {
            self.visit_expr(&expr);

            self.expr_type
        }).unwrap_or(BuiltinType::Unit);

        let fct = self.fct.unwrap();
        let types = self.ctxt.types.borrow();

        let fct_type = *types.get(&fct.id).unwrap();

        if expr_type != fct_type {
            let msg = Msg::ReturnType(fct_type.to_string(), expr_type.to_string());
            self.ctxt.diag.borrow_mut().report(s.pos, msg);
        }
    }

    fn check_expr_ident(&mut self, e: &'ast ExprIdentType) {
        let defs = self.ctxt.defs.borrow();
        let var_id = *defs.get(&e.id).unwrap();

        let types = self.ctxt.types.borrow();
        self.expr_type = *types.get(&var_id).unwrap();
    }

    fn check_expr_assign(&mut self, e: &'ast ExprAssignType) {
        if !e.lhs.is_ident() {
            self.ctxt.diag.borrow_mut().report(e.pos, Msg::LvalueExpected);
            return;
        }

        self.visit_expr(&e.lhs);
        let lhs_type = self.expr_type;

        self.visit_expr(&e.rhs);
        let rhs_type = self.expr_type;

        if lhs_type != rhs_type {
            let lhs_type = lhs_type.to_string();
            let rhs_type = rhs_type.to_string();
            let msg = Msg::AssignType(lhs_type, rhs_type);

            self.ctxt.diag.borrow_mut().report(e.pos, msg);
        }

        self.expr_type = lhs_type;
    }

    fn check_expr_un(&mut self, e: &'ast ExprUnType) {
        let expected_type = if e.op == UnOp::Not {
            BuiltinType::Bool
        } else {
            BuiltinType::Int
        };

        self.visit_expr(&e.opnd);
        let opnd_type = self.expr_type;

        if expected_type != opnd_type {
            let opnd_type = opnd_type.to_string();
            let op = e.op.to_string();
            let msg = Msg::UnOpType(op, opnd_type);

            self.ctxt.diag.borrow_mut().report(e.pos, msg);
        }
    }

    fn check_expr_bin(&mut self, e: &'ast ExprBinType) {
        let expected_type = if e.op == BinOp::Or || e.op == BinOp::And {
            BuiltinType::Bool
        } else {
            BuiltinType::Int
        };

        self.visit_expr(&e.lhs);
        let lhs_type = self.expr_type;

        self.visit_expr(&e.rhs);
        let rhs_type = self.expr_type;

        if expected_type != lhs_type || expected_type != rhs_type {
            let lhs_type = lhs_type.to_string();
            let rhs_type = rhs_type.to_string();
            let op = e.op.to_string();
            let msg = Msg::BinOpType(op, lhs_type, rhs_type);

            self.ctxt.diag.borrow_mut().report(e.pos, msg);
        }
    }
}

impl<'a, 'ast> Visitor<'ast> for TypeCheck<'a, 'ast> {
    fn visit_fct(&mut self, f: &'ast Function) {
        self.fct = Some(f);

        visit::walk_fct(self, f);
    }

    fn visit_expr(&mut self, e: &'ast Expr) {
        match *e {
            ExprLitInt(_) => self.expr_type = BuiltinType::Int,
            ExprLitStr(_) => self.expr_type = BuiltinType::Str,
            ExprLitBool(_) => self.expr_type = BuiltinType::Bool,
            ExprIdent(ref expr) => self.check_expr_ident(expr),
            ExprAssign(ref expr) => self.check_expr_assign(expr),
            ExprUn(ref expr) => self.check_expr_un(expr),
            ExprBin(ref expr) => self.check_expr_bin(expr),

            // TODO: rest of possible expressions
            _ => {
                self.expr_type = BuiltinType::Unit;

                visit::walk_expr(self, e)
            }
        }
    }

    fn visit_stmt(&mut self, s: &'ast Stmt) {
        match *s {
            StmtVar(ref stmt) => self.check_stmt_var(stmt),
            StmtWhile(ref stmt) => self.check_stmt_while(stmt),
            StmtIf(ref stmt) => self.check_stmt_if(stmt),
            StmtReturn(ref stmt) => self.check_stmt_return(stmt),

            // for the rest of the statements, no special handling is necessary
            StmtBreak(ref stmt) => visit::walk_stmt(self, s),
            StmtContinue(ref stmt) => visit::walk_stmt(self, s),
            StmtLoop(ref stmt) => visit::walk_stmt(self, s),
            StmtExpr(ref stmt) => visit::walk_stmt(self, s),
            StmtBlock(ref stmt) => visit::walk_stmt(self, s),
        }
    }
}

#[cfg(test)]
mod tests {
    use error::msg::Msg;
    use semck::tests::*;

    #[test]
    fn type_def_for_return_type() {
        ok("fn a() -> int {}");
        err("fn a() -> unknown {}", pos(1, 11), Msg::UnknownType("unknown".into()));
    }

    #[test]
    fn type_def_for_param() {
        ok("fn a(b: int) {}");
        err("fn a(b: foo) {}", pos(1, 9), Msg::UnknownType("foo".into()));
    }

    #[test]
    fn type_def_for_var() {
        ok("fn a() { var a : int = 1; }");
        err("fn a() { var a : test = 1; }", pos(1, 18), Msg::UnknownType("test".into()));
    }

    #[test]
    fn type_var_needs_expr_or_definition() {
        err("fn a() { var a; }", pos(1, 10), Msg::VarNeedsTypeInfo("a".into()));
    }

    #[test]
    fn type_var_wrong_type_defined() {
        ok("fn f() { var a : int = 1; }");
        ok("fn f() { var a : bool = false; }");
        ok("fn f() { var a : str = \"f\"; }");

        err("fn f() { var a : int = true; }",
            pos(1, 10), Msg::VarTypesIncompatible("a".into(), "int".into(), "bool".into()));
        err("fn f() { var b : bool = 2; }",
            pos(1, 10), Msg::VarTypesIncompatible("b".into(), "bool".into(), "int".into()));
    }

    #[test]
    fn type_while() {
        ok("fn x() { while true { } }");
        ok("fn x() { while false { } }");
        err("fn x() { while 2 { } }", pos(1, 10), Msg::WhileCondType("int".into()));
    }

    #[test]
    fn type_if() {
        ok("fn x() { if true { } }");
        ok("fn x() { if false { } }");
        err("fn x() { if 4 { } }", pos(1, 10), Msg::IfCondType("int".into()));
    }

    #[test]
    fn type_return_unit() {
        ok("fn f() { return; }");
        err("fn f() { return 1; }", pos(1, 10), Msg::ReturnType("()".into(), "int".into()));
    }

    #[test]
    fn type_return() {
        ok("fn f() -> int { return 1; }");
        err("fn f() -> int { return; }", pos(1, 17), Msg::ReturnType("int".into(), "()".into()));
    }

    #[test]
    fn type_variable() {
        ok("fn f(a: int) { var b: int = a; }");
    }

    #[test]
    fn type_assign() {
        ok("fn f(a: int) { a = 1; }");
        err("fn f(a: int) { a = true; }",
            pos(1, 18), Msg::AssignType("int".into(), "bool".into()));
    }

    #[test]
    fn type_assign_lvalue() {
        err("fn f() { 1 = 3; }", pos(1, 12), Msg::LvalueExpected);
    }

    #[test]
    fn type_un_op() {
        ok("fn f(a: int) { ~a; -a; +a; }");
        err("fn f(a: int) { !a; }", pos(1, 16),
            Msg::UnOpType("!".into(), "int".into()));

        err("fn f(a: bool) { ~a; }", pos(1, 17),
            Msg::UnOpType("~".into(), "bool".into()));
        err("fn f(a: bool) { -a; }", pos(1, 17),
            Msg::UnOpType("-".into(), "bool".into()));
        err("fn f(a: bool) { +a; }", pos(1, 17),
            Msg::UnOpType("+".into(), "bool".into()));
    }

    #[test]
    fn type_bin_op() {
        ok("fn f(a: int) { a+a; a-a; a*a; a/a; a%a; }");
        ok("fn f(a: int) { a<a; a<=a; a==a; a!=a; a>a; a>=a; }");
        ok("fn f(a: int) { a|a; a&a; a^a; }");
        ok("fn f(a: bool) { a||a; a&&a; }");

        err("fn f(a: bool) { a+a; }", pos(1, 18),
            Msg::BinOpType("+".into(), "bool".into(), "bool".into()));
        err("fn f(a: bool) { a^a; }", pos(1, 18),
            Msg::BinOpType("^".into(), "bool".into(), "bool".into()));
        err("fn f(a: int) { a||a; }", pos(1, 17),
            Msg::BinOpType("||".into(), "int".into(), "int".into()));
        err("fn f(a: int) { a&&a; }", pos(1, 17),
            Msg::BinOpType("&&".into(), "int".into(), "int".into()));
    }
}
<commit_msg>typeck: check function return type<commit_after>use driver::ctxt::Context;
use error::msg::Msg;

use parser::ast::*;
use parser::ast::Expr::*;
use parser::ast::Stmt::*;
use parser::ast::Type::*;
use parser::ast::visit::Visitor;

use sym::*;
use sym::Sym::*;

pub fn check(ctxt: &Context, ast: &Ast) {
    DefCheck::new(ctxt).visit_ast(ast);

    if ctxt.diag.borrow().has_errors() { return; }

    TypeCheck::new(ctxt).visit_ast(ast);
}

struct DefCheck<'a, 'ast: 'a> {
    ctxt: &'a Context<'a, 'ast>,
    current_id: Option<NodeId>,
}

impl<'a, 'ast> DefCheck<'a, 'ast> {
    fn new(ctxt: &'a Context<'a, 'ast>) -> DefCheck<'a, 'ast> {
        DefCheck {
            ctxt: ctxt,
            current_id: None
        }
    }
}

impl<'a, 'ast> Visitor<'ast> for DefCheck<'a, 'ast> {
    fn visit_fct(&mut self, f: &'ast Function) {
        self.current_id = Some(f.id);

        visit::walk_fct(self, f);

        // if no type for function found, then function returns ()
        self.ctxt.types.borrow_mut().entry(f.id).or_insert(BuiltinType::Unit);
    }

    fn visit_param(&mut self, p: &'ast Param) {
        self.current_id = Some(p.id);

        visit::walk_param(self, p);
    }

    fn visit_stmt(&mut self, s: &'ast Stmt) {
        if let StmtVar(ref var) = *s {
            self.current_id = Some(var.id)
        }

        visit::walk_stmt(self, s);
    }

    fn visit_type(&mut self, t: &'ast Type) {
        match *t {
            TypeBasic(ref basic) => {
                if let Some(builtin) = self.ctxt.sym.borrow().get_type(basic.name) {
                    let id = self.current_id.unwrap();
                    self.ctxt.types.borrow_mut().insert(id, builtin);
                } else {
                    let tyname = self.ctxt.interner.str(basic.name).to_string();
                    let msg = Msg::UnknownType(tyname);
                    self.ctxt.diag.borrow_mut().report(basic.pos, msg);
                }
            }

            _ => self.ctxt.diag.borrow_mut().report_unimplemented(t.pos())
        }
    }
}

struct TypeCheck<'a, 'ast: 'a> {
    ctxt: &'a Context<'a, 'ast>,
    expr_type: BuiltinType,
    fct: Option<&'ast Function>,
}

impl<'a, 'ast> TypeCheck<'a, 'ast> {
    fn new(ctxt: &'a Context<'a, 'ast>) -> TypeCheck<'a, 'ast> {
        TypeCheck {
            ctxt: ctxt,
            expr_type: BuiltinType::Unit,
            fct: None
        }
    }

    fn check_stmt_var(&mut self, s: &'ast StmtVarType) {
        let expr_type = s.expr.as_ref().map(|expr| {
            self.visit_expr(&expr);

            self.expr_type
        });

        let defined_type = if let Some(ref ty) = s.data_type {
            self.ctxt.types.borrow().get(&s.id).map(|bt| *bt)
        } else {
            expr_type
        };

        if defined_type.is_none() {
            let tyname = self.ctxt.interner.str(s.name).to_string();
            self.ctxt.diag.borrow_mut().report(s.pos, Msg::VarNeedsTypeInfo(tyname));

            return;
        }

        if expr_type.is_some() && (defined_type.unwrap() != expr_type.unwrap()) {
            let varname = self.ctxt.interner.str(s.name).to_string();
            let defname = defined_type.unwrap().to_string();
            let exprname = expr_type.unwrap().to_string();
            let msg = Msg::VarTypesIncompatible(varname, defname, exprname);

            self.ctxt.diag.borrow_mut().report(s.pos, msg);
        }
    }

    fn check_stmt_while(&mut self, s: &'ast StmtWhileType) {
        self.visit_expr(&s.cond);

        if self.expr_type != BuiltinType::Bool {
            let tyname = self.expr_type.to_string();
            let msg = Msg::WhileCondType(tyname);

            self.ctxt.diag.borrow_mut().report(s.pos, msg);
        }

        self.visit_stmt(&s.block);
    }

    fn check_stmt_if(&mut self, s: &'ast StmtIfType) {
        self.visit_expr(&s.cond);

        if self.expr_type != BuiltinType::Bool {
            let tyname = self.expr_type.to_string();
            let msg = Msg::IfCondType(tyname);

            self.ctxt.diag.borrow_mut().report(s.pos, msg);
        }

        self.visit_stmt(&s.then_block);

        if let Some(ref else_block) = s.else_block {
            self.visit_stmt(else_block);
        }
    }

    fn check_stmt_return(&mut self, s: &'ast StmtReturnType) {
        let expr_type = s.expr.as_ref().map(|expr| {
            self.visit_expr(&expr);

            self.expr_type
        }).unwrap_or(BuiltinType::Unit);

        let fct = self.fct.unwrap();
        let types = self.ctxt.types.borrow();

        let fct_type = *types.get(&fct.id).unwrap();

        if expr_type != fct_type {
            let msg = Msg::ReturnType(fct_type.to_string(), expr_type.to_string());
            self.ctxt.diag.borrow_mut().report(s.pos, msg);
        }
    }

    fn check_expr_ident(&mut self, e: &'ast ExprIdentType) {
        let defs = self.ctxt.defs.borrow();
        let var_id = *defs.get(&e.id).unwrap();

        let types = self.ctxt.types.borrow();
        self.expr_type = *types.get(&var_id).unwrap();
    }

    fn check_expr_assign(&mut self, e: &'ast ExprAssignType) {
        if !e.lhs.is_ident() {
            self.ctxt.diag.borrow_mut().report(e.pos, Msg::LvalueExpected);
            return;
        }

        self.visit_expr(&e.lhs);
        let lhs_type = self.expr_type;

        self.visit_expr(&e.rhs);
        let rhs_type = self.expr_type;

        if lhs_type != rhs_type {
            let lhs_type = lhs_type.to_string();
            let rhs_type = rhs_type.to_string();
            let msg = Msg::AssignType(lhs_type, rhs_type);

            self.ctxt.diag.borrow_mut().report(e.pos, msg);
        }

        self.expr_type = lhs_type;
    }

    fn check_expr_un(&mut self, e: &'ast ExprUnType) {
        let expected_type = if e.op == UnOp::Not {
            BuiltinType::Bool
        } else {
            BuiltinType::Int
        };

        self.visit_expr(&e.opnd);
        let opnd_type = self.expr_type;

        if expected_type != opnd_type {
            let opnd_type = opnd_type.to_string();
            let op = e.op.to_string();
            let msg = Msg::UnOpType(op, opnd_type);

            self.ctxt.diag.borrow_mut().report(e.pos, msg);
        }
    }

    fn check_expr_bin(&mut self, e: &'ast ExprBinType) {
        let expected_type = if e.op == BinOp::Or || e.op == BinOp::And {
            BuiltinType::Bool
        } else {
            BuiltinType::Int
        };

        self.visit_expr(&e.lhs);
        let lhs_type = self.expr_type;

        self.visit_expr(&e.rhs);
        let rhs_type = self.expr_type;

        if expected_type != lhs_type || expected_type != rhs_type {
            let lhs_type = lhs_type.to_string();
            let rhs_type = rhs_type.to_string();
            let op = e.op.to_string();
            let msg = Msg::BinOpType(op, lhs_type, rhs_type);

            self.ctxt.diag.borrow_mut().report(e.pos, msg);
        }
    }

    fn check_expr_call(&mut self, e: &'ast ExprCallType) {
        let defs = self.ctxt.defs.borrow();
        let var_id = *defs.get(&e.id).unwrap();

        let types = self.ctxt.types.borrow();
        self.expr_type = *types.get(&var_id).unwrap();
    }
}

impl<'a, 'ast> Visitor<'ast> for TypeCheck<'a, 'ast> {
    fn visit_fct(&mut self, f: &'ast Function) {
        self.fct = Some(f);

        visit::walk_fct(self, f);
    }

    fn visit_expr(&mut self, e: &'ast Expr) {
        match *e {
            ExprLitInt(_) => self.expr_type = BuiltinType::Int,
            ExprLitStr(_) => self.expr_type = BuiltinType::Str,
            ExprLitBool(_) => self.expr_type = BuiltinType::Bool,
            ExprIdent(ref expr) => self.check_expr_ident(expr),
            ExprAssign(ref expr) => self.check_expr_assign(expr),
            ExprUn(ref expr) => self.check_expr_un(expr),
            ExprBin(ref expr) => self.check_expr_bin(expr),
            ExprCall(ref expr) => self.check_expr_call(expr),
        }
    }

    fn visit_stmt(&mut self, s: &'ast Stmt) {
        match *s {
            StmtVar(ref stmt) => self.check_stmt_var(stmt),
            StmtWhile(ref stmt) => self.check_stmt_while(stmt),
            StmtIf(ref stmt) => self.check_stmt_if(stmt),
            StmtReturn(ref stmt) => self.check_stmt_return(stmt),

            // for the rest of the statements, no special handling is necessary
            StmtBreak(ref stmt) => visit::walk_stmt(self, s),
            StmtContinue(ref stmt) => visit::walk_stmt(self, s),
            StmtLoop(ref stmt) => visit::walk_stmt(self, s),
            StmtExpr(ref stmt) => visit::walk_stmt(self, s),
            StmtBlock(ref stmt) => visit::walk_stmt(self, s),
        }
    }
}

#[cfg(test)]
mod tests {
    use error::msg::Msg;
    use semck::tests::*;

    #[test]
    fn type_def_for_return_type() {
        ok("fn a() -> int {}");
        err("fn a() -> unknown {}", pos(1, 11), Msg::UnknownType("unknown".into()));
    }

    #[test]
    fn type_def_for_param() {
        ok("fn a(b: int) {}");
        err("fn a(b: foo) {}", pos(1, 9), Msg::UnknownType("foo".into()));
    }

    #[test]
    fn type_def_for_var() {
        ok("fn a() { var a : int = 1; }");
        err("fn a() { var a : test = 1; }", pos(1, 18), Msg::UnknownType("test".into()));
    }

    #[test]
    fn type_var_needs_expr_or_definition() {
        err("fn a() { var a; }", pos(1, 10), Msg::VarNeedsTypeInfo("a".into()));
    }

    #[test]
    fn type_var_wrong_type_defined() {
        ok("fn f() { var a : int = 1; }");
        ok("fn f() { var a : bool = false; }");
        ok("fn f() { var a : str = \"f\"; }");

        err("fn f() { var a : int = true; }",
            pos(1, 10), Msg::VarTypesIncompatible("a".into(), "int".into(), "bool".into()));
        err("fn f() { var b : bool = 2; }",
            pos(1, 10), Msg::VarTypesIncompatible("b".into(), "bool".into(), "int".into()));
    }

    #[test]
    fn type_while() {
        ok("fn x() { while true { } }");
        ok("fn x() { while false { } }");
        err("fn x() { while 2 { } }", pos(1, 10), Msg::WhileCondType("int".into()));
    }

    #[test]
    fn type_if() {
        ok("fn x() { if true { } }");
        ok("fn x() { if false { } }");
        err("fn x() { if 4 { } }", pos(1, 10), Msg::IfCondType("int".into()));
    }

    #[test]
    fn type_return_unit() {
        ok("fn f() { return; }");
        err("fn f() { return 1; }", pos(1, 10), Msg::ReturnType("()".into(), "int".into()));
    }

    #[test]
    fn type_return() {
        ok("fn f() -> int { return 1; }");
        err("fn f() -> int { return; }", pos(1, 17), Msg::ReturnType("int".into(), "()".into()));
    }

    #[test]
    fn type_variable() {
        ok("fn f(a: int) { var b: int = a; }");
    }

    #[test]
    fn type_assign() {
        ok("fn f(a: int) { a = 1; }");
        err("fn f(a: int) { a = true; }",
            pos(1, 18), Msg::AssignType("int".into(), "bool".into()));
    }

    #[test]
    fn type_assign_lvalue() {
        err("fn f() { 1 = 3; }", pos(1, 12), Msg::LvalueExpected);
    }

    #[test]
    fn type_un_op() {
        ok("fn f(a: int) { ~a; -a; +a; }");
        err("fn f(a: int) { !a; }", pos(1, 16),
            Msg::UnOpType("!".into(), "int".into()));

        err("fn f(a: bool) { ~a; }", pos(1, 17),
            Msg::UnOpType("~".into(), "bool".into()));
        err("fn f(a: bool) { -a; }", pos(1, 17),
            Msg::UnOpType("-".into(), "bool".into()));
        err("fn f(a: bool) { +a; }", pos(1, 17),
            Msg::UnOpType("+".into(), "bool".into()));
    }

    #[test]
    fn type_bin_op() {
        ok("fn f(a: int) { a+a; a-a; a*a; a/a; a%a; }");
        ok("fn f(a: int) { a<a; a<=a; a==a; a!=a; a>a; a>=a; }");
        ok("fn f(a: int) { a|a; a&a; a^a; }");
        ok("fn f(a: bool) { a||a; a&&a; }");

        err("fn f(a: bool) { a+a; }", pos(1, 18),
            Msg::BinOpType("+".into(), "bool".into(), "bool".into()));
        err("fn f(a: bool) { a^a; }", pos(1, 18),
            Msg::BinOpType("^".into(), "bool".into(), "bool".into()));
        err("fn f(a: int) { a||a; }", pos(1, 17),
            Msg::BinOpType("||".into(), "int".into(), "int".into()));
        err("fn f(a: int) { a&&a; }", pos(1, 17),
            Msg::BinOpType("&&".into(), "int".into(), "int".into()));
    }

    #[test]
    fn type_function_return_type() {
        ok("fn foo() -> int { return 1; }\nfn f() { var i: int = foo(); }");
        err("fn foo() -> int { return 1; }\nfn f() { var i: bool = foo(); }",
            pos(2, 10), Msg::VarTypesIncompatible("i".into(), "bool".into(), "int".into()));
    }
}
<|endoftext|>


<commit_before>use std::path::PathBuf;

use glob::glob;
use glob::Paths;
use glob::PatternError;

use storage::file::id::FileID;
use storage::file::id_type::FileIDType;
use storage::file::hash::FileHash;
use module::Module;

/*
 * A path represents either a GLOB ("/tmp/store/module-*-*.imag" for example) or a full path
 *
 * It can be used to generate a File or iterate over some files
 *
 */
struct Path<'a> {

    /*
     * The base part ("/tmp/")
     */
    base: PathBuf,

    /*
     * The store part ("/store/")
     */
    store: PathBuf,

    /*
     * The module
     */
    module: &'a Module<'a>,

    /*
     * The ID
     */
    idtype: Option<FileIDType>,
    idhash: Option<FileHash>,
    id:     Option<FileID>,

}

impl<'a> Path<'a> {

    fn new(base: PathBuf, store: PathBuf, m: &'a Module<'a>, id: FileID) -> Path<'a> {
        Path {
            base:   base,
            store:  store,
            module: m,
            idtype: Some(id.get_type()),
            idhash: Some(id.get_id()),
            id:     Some(id),
        }
    }

    fn new_with_idtype(base: PathBuf, store: PathBuf, m: &'a Module<'a>, id: FileIDType) -> Path<'a> {
        Path {
            base:   base,
            store:  store,
            module: m,
            idtype: Some(id),
            idhash: None,
            id:     None,
        }
    }

    fn new_with_idhash(base: PathBuf, store: PathBuf, m: &'a Module<'a>, id: FileHash) -> Path<'a> {
        Path {
            base:   base,
            store:  store,
            module: m,
            idtype: None,
            idhash: Some(id),
            id:     None,
        }
    }

}

/*
 * Transform Path into str, so we can call glob() on it
 */
impl<'a> Into<&'a str> for Path<'a> {

    fn into(self) -> &'a str {
        let s = self.base.clone();
        s.push(self.store.clone());
        s.push(self.module.name());
        if self.id.is_some() {
            let idstr : String = self.id.unwrap().into();
            s.push(idstr);
        } else {
            s.push("*");
        }
        s.set_extension("imag");
        s.to_str().unwrap_or("")
    }
}

impl<'a> Into<Result<Paths, PatternError>> for Path<'a> {

    fn into(self) -> Result<Paths, PatternError> {
        let s : &str = self.into();
        glob(s)
    }

}
<commit_msg>Fixup Into<str> for Path<commit_after>use std::path::PathBuf;

use glob::glob;
use glob::Paths;
use glob::PatternError;

use storage::file::id::FileID;
use storage::file::id_type::FileIDType;
use storage::file::hash::FileHash;
use module::Module;

/*
 * A path represents either a GLOB ("/tmp/store/module-*-*.imag" for example) or a full path
 *
 * It can be used to generate a File or iterate over some files
 *
 */
struct Path<'a> {

    /*
     * The base part ("/tmp/")
     */
    base: PathBuf,

    /*
     * The store part ("/store/")
     */
    store: PathBuf,

    /*
     * The module
     */
    module: &'a Module<'a>,

    /*
     * The ID
     */
    idtype: Option<FileIDType>,
    idhash: Option<FileHash>,
    id:     Option<FileID>,

}

impl<'a> Path<'a> {

    fn new(base: PathBuf, store: PathBuf, m: &'a Module<'a>, id: FileID) -> Path<'a> {
        Path {
            base:   base,
            store:  store,
            module: m,
            idtype: Some(id.get_type()),
            idhash: Some(id.get_id()),
            id:     Some(id),
        }
    }

    fn new_with_idtype(base: PathBuf, store: PathBuf, m: &'a Module<'a>, id: FileIDType) -> Path<'a> {
        Path {
            base:   base,
            store:  store,
            module: m,
            idtype: Some(id),
            idhash: None,
            id:     None,
        }
    }

    fn new_with_idhash(base: PathBuf, store: PathBuf, m: &'a Module<'a>, id: FileHash) -> Path<'a> {
        Path {
            base:   base,
            store:  store,
            module: m,
            idtype: None,
            idhash: Some(id),
            id:     None,
        }
    }

}

/*
 * Transform Path into str, so we can call glob() on it
 */
impl<'a> Into<&'a str> for Path<'a> {

    fn into(self) -> &'a str {
        let mut s = self.base.clone();
        s.push(self.store.clone());
        s.push(self.module.name());
        if self.id.is_some() {
            let idstr : String = self.id.unwrap().into();
            s.push(idstr);
        } else {
            s.push("*");
        }
        s.set_extension("imag");
        s.to_str().unwrap_or("")
    }
}

impl<'a> Into<Result<Paths, PatternError>> for Path<'a> {

    fn into(self) -> Result<Paths, PatternError> {
        let s : &str = self.into();
        glob(s)
    }

}
<|endoftext|>


<commit_before>//
//   Copyright 2016 Andrew Hunter
//
//   Licensed under the Apache License, Version 2.0 (the "License");
//   you may not use this file except in compliance with the License.
//   You may obtain a copy of the License at
//
//       http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.
//

//!
//! # Symbol range
//!
//! It's quite common to want to match ranges of symbols (most particularly, 'all' symbols but also things like all upper-case
//! characters). It would be inefficient to store all of the transitions represented by one of these transitions on a per-symbol
//! basis, and impossible for some symbol sets (consider a state machine working on u32s). Instead, transitions are stored as
//! symbol ranges.
//!

///
/// Trait implemented by symbol types that are countable - ie, for types where there's always a next symbol
///
pub trait Countable where Self: Sized {
    ///
    /// Returns the next symbol in order (None if this is the last symbol)
    ///
    fn next(&self) -> Option<Self>;

    ///
    /// Returns the next symbol in order (None if this is the last symbol)
    ///
    fn prev(&self) -> Option<Self>;
}

///
/// Represents a range of symbols
///
/// Symbols must be ordered in order to use them with a range-based state machine.
///
#[derive(Clone, Eq, PartialEq)]
pub struct SymbolRange<Symbol: Ord+Clone> {
    ///
    /// Lowest symbol in the range
    ///
    pub lowest: Symbol,

    ///
    /// Highest symbol in the range
    ///
    /// This is inclusive, so the highest symbol is always included in the range (this makes it differ from Rust's
    /// built-in Range struct, and is important for supporting uncountable symbols)
    ///
    pub highest: Symbol
}

impl<Symbol: Ord+Clone> SymbolRange<Symbol> {
    ///
    /// Creates a new range covering everything between the specified two symbols
    ///
    #[inline]
    pub fn new(lowest: Symbol, highest: Symbol) -> SymbolRange<Symbol> {
        if lowest > highest {
            SymbolRange { lowest: highest, highest: lowest }
        } else {
            SymbolRange { lowest: lowest, highest: highest }
        }
    }

    ///
    /// True if this range overlaps another
    ///
    #[inline]
    pub fn overlaps(&self, with: &SymbolRange<Symbol>) -> bool {
        if self.highest < with.lowest {
            false
        } else if self.lowest > with.highest {
            false
        } else {
            true
        }
    }
}

impl Countable for u8 {
    fn next(&self) -> Option<u8> {
        if *self == Self::max_value() {
            None
        } else {
            Some(self+1)
        }
    }

    fn prev(&self) -> Option<u8> {
        if *self == Self::min_value() {
            None
        } else {
            Some(self-1)
        }
    }
}

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

    #[test]
    fn can_create_range() {
        let range = SymbolRange::new(1, 2);

        assert!(range.lowest == 1);
        assert!(range.highest == 2);
    }

    #[test]
    fn can_create_range_reversed() {
        let range = SymbolRange::new(5, 1);

        assert!(range.lowest == 1);
        assert!(range.highest == 5);
    }

    #[test]
    fn overlaps_when_within() {
        assert!(SymbolRange::new(1, 4).overlaps(&SymbolRange::new(2, 3)));
    }

    #[test]
    fn overlaps_when_without() {
        assert!(SymbolRange::new(2, 3).overlaps(&SymbolRange::new(1, 4)));
    }

    #[test]
    fn overlaps_when_lower() {
        assert!(SymbolRange::new(1, 3).overlaps(&SymbolRange::new(2, 4)));
    }

    #[test]
    fn overlaps_when_higher() {
        assert!(SymbolRange::new(2, 4).overlaps(&SymbolRange::new(1, 3)));
    }

    #[test]
    fn overlaps_when_same() {
        assert!(SymbolRange::new(1, 4).overlaps(&SymbolRange::new(1, 4)));
    }

    #[test]
    fn does_not_overlap_lower() {
        assert!(!SymbolRange::new(1, 2).overlaps(&SymbolRange::new(4, 5)));
    }

    #[test]
    fn does_not_overlap_higher() {
        assert!(!SymbolRange::new(4, 5).overlaps(&SymbolRange::new(1, 2)));
    }
}
<commit_msg>Add a way to join two symbol ranges into one<commit_after>//
//   Copyright 2016 Andrew Hunter
//
//   Licensed under the Apache License, Version 2.0 (the "License");
//   you may not use this file except in compliance with the License.
//   You may obtain a copy of the License at
//
//       http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.
//

//!
//! # Symbol range
//!
//! It's quite common to want to match ranges of symbols (most particularly, 'all' symbols but also things like all upper-case
//! characters). It would be inefficient to store all of the transitions represented by one of these transitions on a per-symbol
//! basis, and impossible for some symbol sets (consider a state machine working on u32s). Instead, transitions are stored as
//! symbol ranges.
//!

///
/// Trait implemented by symbol types that are countable - ie, for types where there's always a next symbol
///
pub trait Countable where Self: Sized {
    ///
    /// Returns the next symbol in order (None if this is the last symbol)
    ///
    fn next(&self) -> Option<Self>;

    ///
    /// Returns the next symbol in order (None if this is the last symbol)
    ///
    fn prev(&self) -> Option<Self>;
}

///
/// Represents a range of symbols
///
/// Symbols must be ordered in order to use them with a range-based state machine.
///
#[derive(Clone, Eq, PartialEq)]
pub struct SymbolRange<Symbol: Ord+Clone> {
    ///
    /// Lowest symbol in the range
    ///
    pub lowest: Symbol,

    ///
    /// Highest symbol in the range
    ///
    /// This is inclusive, so the highest symbol is always included in the range (this makes it differ from Rust's
    /// built-in Range struct, and is important for supporting uncountable symbols)
    ///
    pub highest: Symbol
}

impl<Symbol: Ord+Clone> SymbolRange<Symbol> {
    ///
    /// Creates a new range covering everything between the specified two symbols
    ///
    #[inline]
    pub fn new(lowest: Symbol, highest: Symbol) -> SymbolRange<Symbol> {
        if lowest > highest {
            SymbolRange { lowest: highest, highest: lowest }
        } else {
            SymbolRange { lowest: lowest, highest: highest }
        }
    }

    ///
    /// Joins this range with another
    ///
    /// This creates a new range that covers all the symbols of both. If `overlaps()` is false for these two ranges, then
    /// the new range may cover additional symbols that are not in either range.
    ///
    pub fn join(&self, with: &SymbolRange<Symbol>) -> SymbolRange<Symbol> {
        SymbolRange {
            lowest:  if with.lowest<self.lowest   { with.lowest.clone()  } else { self.lowest.clone()  },
            highest: if with.highest<self.highest { self.highest.clone() } else { with.highest.clone() }
        }
    }

    ///
    /// True if this range overlaps another
    ///
    #[inline]
    pub fn overlaps(&self, with: &SymbolRange<Symbol>) -> bool {
        if self.highest < with.lowest {
            false
        } else if self.lowest > with.highest {
            false
        } else {
            true
        }
    }
}

impl Countable for u8 {
    fn next(&self) -> Option<u8> {
        if *self == Self::max_value() {
            None
        } else {
            Some(self+1)
        }
    }

    fn prev(&self) -> Option<u8> {
        if *self == Self::min_value() {
            None
        } else {
            Some(self-1)
        }
    }
}

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

    #[test]
    fn can_create_range() {
        let range = SymbolRange::new(1, 2);

        assert!(range.lowest == 1);
        assert!(range.highest == 2);
    }

    #[test]
    fn can_create_range_reversed() {
        let range = SymbolRange::new(5, 1);

        assert!(range.lowest == 1);
        assert!(range.highest == 5);
    }

    #[test]
    fn overlaps_when_within() {
        assert!(SymbolRange::new(1, 4).overlaps(&SymbolRange::new(2, 3)));
    }

    #[test]
    fn overlaps_when_without() {
        assert!(SymbolRange::new(2, 3).overlaps(&SymbolRange::new(1, 4)));
    }

    #[test]
    fn overlaps_when_lower() {
        assert!(SymbolRange::new(1, 3).overlaps(&SymbolRange::new(2, 4)));
    }

    #[test]
    fn overlaps_when_higher() {
        assert!(SymbolRange::new(2, 4).overlaps(&SymbolRange::new(1, 3)));
    }

    #[test]
    fn overlaps_when_same() {
        assert!(SymbolRange::new(1, 4).overlaps(&SymbolRange::new(1, 4)));
    }

    #[test]
    fn does_not_overlap_lower() {
        assert!(!SymbolRange::new(1, 2).overlaps(&SymbolRange::new(4, 5)));
    }

    #[test]
    fn does_not_overlap_higher() {
        assert!(!SymbolRange::new(4, 5).overlaps(&SymbolRange::new(1, 2)));
    }

    #[test]
    fn join_left() {
        let joined = SymbolRange::new(1, 3).join(&SymbolRange::new(3, 4));

        assert!(joined.lowest == 1);
        assert!(joined.highest == 4);
    }

    #[test]
    fn join_right() {
        let joined = SymbolRange::new(3, 4).join(&SymbolRange::new(1, 3));

        assert!(joined.lowest == 1);
        assert!(joined.highest == 4);
    }

    #[test]
    fn join_overlap() {
        let joined = SymbolRange::new(1, 4).join(&SymbolRange::new(2, 3));

        assert!(joined.lowest == 1);
        assert!(joined.highest == 4);
    }
}
<|endoftext|>


<commit_before>// Copyright 2012-2014 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.

//! Terminfo database interface.

use std::collections::HashMap;
use std::env;
use std::fs::File;
use std::io::prelude::*;
use std::io;
use std::io::BufReader;
use std::path::Path;

use Attr;
use color;
use Terminal;
use Result;
use self::searcher::get_dbpath_for_term;
use self::parser::compiled::{parse, msys_terminfo};
use self::parm::{expand, Variables, Param};
use self::Error::*;

/// A parsed terminfo database entry.
#[derive(Debug)]
pub struct TermInfo {
    /// Names for the terminal
    pub names: Vec<String>,
    /// Map of capability name to boolean value
    pub bools: HashMap<String, bool>,
    /// Map of capability name to numeric value
    pub numbers: HashMap<String, u16>,
    /// Map of capability name to raw (unexpanded) string
    pub strings: HashMap<String, Vec<u8>>,
}

impl TermInfo {
    /// Create a TermInfo based on current environment.
    pub fn from_env() -> Result<TermInfo> {
        let term = match env::var("TERM") {
            Ok(name) => TermInfo::from_name(&name),
            Err(..) => return Err(::Error::TermUnset),
        };

        if term.is_err() && env::var("MSYSCON").ok().map_or(false, |s| "mintty.exe" == s) {
            // msys terminal
            Ok(msys_terminfo())
        } else {
            term
        }
    }

    /// Create a TermInfo for the named terminal.
    pub fn from_name(name: &str) -> Result<TermInfo> {
        get_dbpath_for_term(name)
            .ok_or_else(|| ::Error::TerminfoEntryNotFound)
            .and_then(|p| TermInfo::from_path(&p))
    }

    /// Parse the given TermInfo.
    pub fn from_path<P: AsRef<Path>>(path: P) -> Result<TermInfo> {
        Self::_from_path(path.as_ref())
    }
    // Keep the metadata small
    // (That is, this uses a &Path so that this function need not be instantiated
    // for every type
    // which implements AsRef<Path>. One day, if/when rustc is a bit smarter, it
    // might do this for
    // us. Alas. )
    fn _from_path(path: &Path) -> Result<TermInfo> {
        let file = try!(File::open(path).map_err(|e| ::Error::Io(e)));
        let mut reader = BufReader::new(file);
        parse(&mut reader, false)
    }
}

#[derive(Debug, Eq, PartialEq)]
/// An error from parsing a terminfo entry
pub enum Error {
    /// The "magic" number at the start of the file was wrong.
    ///
    /// It should be `0x11A`
    BadMagic(u16),
    /// The names in the file were not valid UTF-8.
    ///
    /// In theory these should only be ASCII, but to work with the Rust `str` type, we treat them
    /// as UTF-8. This is valid, except when a terminfo file decides to be invalid. This hasn't
    /// been encountered in the wild.
    NotUtf8(::std::str::Utf8Error),
    /// The names section of the file was empty
    ShortNames,
    /// More boolean parameters are present in the file than this crate knows how to interpret.
    TooManyBools,
    /// More number parameters are present in the file than this crate knows how to interpret.
    TooManyNumbers,
    /// More string parameters are present in the file than this crate knows how to interpret.
    TooManyStrings,
    /// The length of some field was not >= -1.
    InvalidLength,
    /// The names table was missing a trailing null terminator.
    NamesMissingNull,
    /// The strings table was missing a trailing null terminator.
    StringsMissingNull,
}

impl ::std::fmt::Display for Error {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        use std::error::Error;
        match self {
            &NotUtf8(e) => write!(f, "{}", e),
            &BadMagic(v) => write!(f, "bad magic number {:x} in terminfo header", v),
            _ => f.write_str(self.description()),
        }
    }
}

impl ::std::convert::From<::std::string::FromUtf8Error> for Error {
    fn from(v: ::std::string::FromUtf8Error) -> Self {
        NotUtf8(v.utf8_error())
    }
}

impl ::std::error::Error for Error {
    fn description(&self) -> &str {
        match self {
            &BadMagic(..) => "incorrect magic number at start of file",
            &ShortNames => "no names exposed, need at least one",
            &TooManyBools => "more boolean properties than libterm knows about",
            &TooManyNumbers => "more number properties than libterm knows about",
            &TooManyStrings => "more string properties than libterm knows about",
            &InvalidLength => "invalid length field value, must be >= -1",
            &NotUtf8(ref e) => e.description(),
            &NamesMissingNull => "names table missing NUL terminator",
            &StringsMissingNull => "string table missing NUL terminator",
        }
    }

    fn cause(&self) -> Option<&::std::error::Error> {
        match self {
            &NotUtf8(ref e) => Some(e),
            _ => None,
        }
    }
}

pub mod searcher;

/// TermInfo format parsing.
pub mod parser {
    //! ncurses-compatible compiled terminfo format parsing (term(5))
    pub mod compiled;
    mod names;
}
pub mod parm;

fn cap_for_attr(attr: Attr) -> &'static str {
    match attr {
        Attr::Bold => "bold",
        Attr::Dim => "dim",
        Attr::Italic(true) => "sitm",
        Attr::Italic(false) => "ritm",
        Attr::Underline(true) => "smul",
        Attr::Underline(false) => "rmul",
        Attr::Blink => "blink",
        Attr::Standout(true) => "smso",
        Attr::Standout(false) => "rmso",
        Attr::Reverse => "rev",
        Attr::Secure => "invis",
        Attr::ForegroundColor(_) => "setaf",
        Attr::BackgroundColor(_) => "setab",
    }
}

/// A Terminal that knows how many colors it supports, with a reference to its
/// parsed Terminfo database record.
pub struct TerminfoTerminal<T> {
    num_colors: u16,
    out: T,
    ti: TermInfo,
}

impl<T: Write + Send> Terminal for TerminfoTerminal<T> {
    type Output = T;
    fn fg(&mut self, color: color::Color) -> Result<()> {
        let color = self.dim_if_necessary(color);
        if self.num_colors > color {
            return self.apply_cap("setaf", &[Param::Number(color as i32)]);
        }
        Err(::Error::ColorOutOfRange)
    }

    fn bg(&mut self, color: color::Color) -> Result<()> {
        let color = self.dim_if_necessary(color);
        if self.num_colors > color {
            return self.apply_cap("setab", &[Param::Number(color as i32)]);
        }
        Err(::Error::ColorOutOfRange)
    }

    fn attr(&mut self, attr: Attr) -> Result<()> {
        match attr {
            Attr::ForegroundColor(c) => self.fg(c),
            Attr::BackgroundColor(c) => self.bg(c),
            _ => self.apply_cap(cap_for_attr(attr), &[]),
        }
    }

    fn supports_attr(&self, attr: Attr) -> bool {
        match attr {
            Attr::ForegroundColor(_) | Attr::BackgroundColor(_) => self.num_colors > 0,
            _ => {
                let cap = cap_for_attr(attr);
                self.ti.strings.get(cap).is_some()
            }
        }
    }

    fn reset(&mut self) -> Result<()> {
        // are there any terminals that have color/attrs and not sg0?
        // Try falling back to sgr, then op
        let cmd = match ["sg0", "sgr", "op"]
                            .iter()
                            .filter_map(|cap| self.ti.strings.get(*cap))
                            .next() {
            Some(op) => {
                match expand(&op, &[], &mut Variables::new()) {
                    Ok(cmd) => cmd,
                    Err(e) => return Err(e.into()),
                }
            }
            None => return Err(::Error::NotSupported),
        };
        try!(self.out.write_all(&cmd));
        Ok(())
    }

    fn cursor_up(&mut self) -> Result<()> {
        self.apply_cap("cuu1", &[])
    }

    fn delete_line(&mut self) -> Result<()> {
        self.apply_cap("dl", &[])
    }

    fn carriage_return(&mut self) -> Result<()> {
        self.apply_cap("cr", &[])
    }

    fn get_ref<'a>(&'a self) -> &'a T {
        &self.out
    }

    fn get_mut<'a>(&'a mut self) -> &'a mut T {
        &mut self.out
    }

    fn into_inner(self) -> T
        where Self: Sized
    {
        self.out
    }
}

impl<T: Write + Send> TerminfoTerminal<T> {
    /// Create a new TerminfoTerminal with the given TermInfo and Write.
    pub fn new_with_terminfo(out: T, terminfo: TermInfo) -> TerminfoTerminal<T> {
        let nc = if terminfo.strings.contains_key("setaf") &&
                    terminfo.strings.contains_key("setab") {
            terminfo.numbers.get("colors").map_or(0, |&n| n)
        } else {
            0
        };

        TerminfoTerminal {
            out: out,
            ti: terminfo,
            num_colors: nc,
        }
    }

    /// Create a new TerminfoTerminal for the current environment with the given Write.
    ///
    /// Returns `None` when the terminfo cannot be found or parsed.
    pub fn new(out: T) -> Option<TerminfoTerminal<T>> {
        TermInfo::from_env().map(move |ti| TerminfoTerminal::new_with_terminfo(out, ti)).ok()
    }

    fn dim_if_necessary(&self, color: color::Color) -> color::Color {
        if color >= self.num_colors && color >= 8 && color < 16 {
            color - 8
        } else {
            color
        }
    }

    fn apply_cap(&mut self, cmd: &str, params: &[Param]) -> Result<()> {
        match self.ti.strings.get(cmd) {
            Some(cmd) => {
                match expand(&cmd, params, &mut Variables::new()) {
                    Ok(s) => {
                        try!(self.out.write_all(&s));
                        Ok(())
                    }
                    Err(e) => Err(e.into()),
                }
            }
            None => Err(::Error::NotSupported),
        }
    }
}

impl<T: Write> Write for TerminfoTerminal<T> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.out.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.out.flush()
    }
}
<commit_msg>Fix reset<commit_after>// Copyright 2012-2014 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.

//! Terminfo database interface.

use std::collections::HashMap;
use std::env;
use std::fs::File;
use std::io::prelude::*;
use std::io;
use std::io::BufReader;
use std::path::Path;

use Attr;
use color;
use Terminal;
use Result;
use self::searcher::get_dbpath_for_term;
use self::parser::compiled::{parse, msys_terminfo};
use self::parm::{expand, Variables, Param};
use self::Error::*;

/// A parsed terminfo database entry.
#[derive(Debug)]
pub struct TermInfo {
    /// Names for the terminal
    pub names: Vec<String>,
    /// Map of capability name to boolean value
    pub bools: HashMap<String, bool>,
    /// Map of capability name to numeric value
    pub numbers: HashMap<String, u16>,
    /// Map of capability name to raw (unexpanded) string
    pub strings: HashMap<String, Vec<u8>>,
}

impl TermInfo {
    /// Create a TermInfo based on current environment.
    pub fn from_env() -> Result<TermInfo> {
        let term = match env::var("TERM") {
            Ok(name) => TermInfo::from_name(&name),
            Err(..) => return Err(::Error::TermUnset),
        };

        if term.is_err() && env::var("MSYSCON").ok().map_or(false, |s| "mintty.exe" == s) {
            // msys terminal
            Ok(msys_terminfo())
        } else {
            term
        }
    }

    /// Create a TermInfo for the named terminal.
    pub fn from_name(name: &str) -> Result<TermInfo> {
        get_dbpath_for_term(name)
            .ok_or_else(|| ::Error::TerminfoEntryNotFound)
            .and_then(|p| TermInfo::from_path(&p))
    }

    /// Parse the given TermInfo.
    pub fn from_path<P: AsRef<Path>>(path: P) -> Result<TermInfo> {
        Self::_from_path(path.as_ref())
    }
    // Keep the metadata small
    // (That is, this uses a &Path so that this function need not be instantiated
    // for every type
    // which implements AsRef<Path>. One day, if/when rustc is a bit smarter, it
    // might do this for
    // us. Alas. )
    fn _from_path(path: &Path) -> Result<TermInfo> {
        let file = try!(File::open(path).map_err(|e| ::Error::Io(e)));
        let mut reader = BufReader::new(file);
        parse(&mut reader, false)
    }
}

#[derive(Debug, Eq, PartialEq)]
/// An error from parsing a terminfo entry
pub enum Error {
    /// The "magic" number at the start of the file was wrong.
    ///
    /// It should be `0x11A`
    BadMagic(u16),
    /// The names in the file were not valid UTF-8.
    ///
    /// In theory these should only be ASCII, but to work with the Rust `str` type, we treat them
    /// as UTF-8. This is valid, except when a terminfo file decides to be invalid. This hasn't
    /// been encountered in the wild.
    NotUtf8(::std::str::Utf8Error),
    /// The names section of the file was empty
    ShortNames,
    /// More boolean parameters are present in the file than this crate knows how to interpret.
    TooManyBools,
    /// More number parameters are present in the file than this crate knows how to interpret.
    TooManyNumbers,
    /// More string parameters are present in the file than this crate knows how to interpret.
    TooManyStrings,
    /// The length of some field was not >= -1.
    InvalidLength,
    /// The names table was missing a trailing null terminator.
    NamesMissingNull,
    /// The strings table was missing a trailing null terminator.
    StringsMissingNull,
}

impl ::std::fmt::Display for Error {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        use std::error::Error;
        match self {
            &NotUtf8(e) => write!(f, "{}", e),
            &BadMagic(v) => write!(f, "bad magic number {:x} in terminfo header", v),
            _ => f.write_str(self.description()),
        }
    }
}

impl ::std::convert::From<::std::string::FromUtf8Error> for Error {
    fn from(v: ::std::string::FromUtf8Error) -> Self {
        NotUtf8(v.utf8_error())
    }
}

impl ::std::error::Error for Error {
    fn description(&self) -> &str {
        match self {
            &BadMagic(..) => "incorrect magic number at start of file",
            &ShortNames => "no names exposed, need at least one",
            &TooManyBools => "more boolean properties than libterm knows about",
            &TooManyNumbers => "more number properties than libterm knows about",
            &TooManyStrings => "more string properties than libterm knows about",
            &InvalidLength => "invalid length field value, must be >= -1",
            &NotUtf8(ref e) => e.description(),
            &NamesMissingNull => "names table missing NUL terminator",
            &StringsMissingNull => "string table missing NUL terminator",
        }
    }

    fn cause(&self) -> Option<&::std::error::Error> {
        match self {
            &NotUtf8(ref e) => Some(e),
            _ => None,
        }
    }
}

pub mod searcher;

/// TermInfo format parsing.
pub mod parser {
    //! ncurses-compatible compiled terminfo format parsing (term(5))
    pub mod compiled;
    mod names;
}
pub mod parm;

fn cap_for_attr(attr: Attr) -> &'static str {
    match attr {
        Attr::Bold => "bold",
        Attr::Dim => "dim",
        Attr::Italic(true) => "sitm",
        Attr::Italic(false) => "ritm",
        Attr::Underline(true) => "smul",
        Attr::Underline(false) => "rmul",
        Attr::Blink => "blink",
        Attr::Standout(true) => "smso",
        Attr::Standout(false) => "rmso",
        Attr::Reverse => "rev",
        Attr::Secure => "invis",
        Attr::ForegroundColor(_) => "setaf",
        Attr::BackgroundColor(_) => "setab",
    }
}

/// A Terminal that knows how many colors it supports, with a reference to its
/// parsed Terminfo database record.
pub struct TerminfoTerminal<T> {
    num_colors: u16,
    out: T,
    ti: TermInfo,
}

impl<T: Write + Send> Terminal for TerminfoTerminal<T> {
    type Output = T;
    fn fg(&mut self, color: color::Color) -> Result<()> {
        let color = self.dim_if_necessary(color);
        if self.num_colors > color {
            return self.apply_cap("setaf", &[Param::Number(color as i32)]);
        }
        Err(::Error::ColorOutOfRange)
    }

    fn bg(&mut self, color: color::Color) -> Result<()> {
        let color = self.dim_if_necessary(color);
        if self.num_colors > color {
            return self.apply_cap("setab", &[Param::Number(color as i32)]);
        }
        Err(::Error::ColorOutOfRange)
    }

    fn attr(&mut self, attr: Attr) -> Result<()> {
        match attr {
            Attr::ForegroundColor(c) => self.fg(c),
            Attr::BackgroundColor(c) => self.bg(c),
            _ => self.apply_cap(cap_for_attr(attr), &[]),
        }
    }

    fn supports_attr(&self, attr: Attr) -> bool {
        match attr {
            Attr::ForegroundColor(_) | Attr::BackgroundColor(_) => self.num_colors > 0,
            _ => {
                let cap = cap_for_attr(attr);
                self.ti.strings.get(cap).is_some()
            }
        }
    }

    fn reset(&mut self) -> Result<()> {
        // are there any terminals that have color/attrs and not sgr0?
        // Try falling back to sgr, then op
        let cmd = match [("sgr0", &[] as &[Param]),
                         ("sgr", &[Param::Number(0)]),
                         ("op", &[])]
                            .iter()
                            .filter_map(|&(cap, params)| self.ti.strings.get(cap).map(|c| (c, params)))
                            .next() {
            Some((op, params)) => {
                match expand(op, params, &mut Variables::new()) {
                    Ok(cmd) => cmd,
                    Err(e) => return Err(e.into()),
                }
            }
            None => return Err(::Error::NotSupported),
        };
        try!(self.out.write_all(&cmd));
        Ok(())
    }

    fn cursor_up(&mut self) -> Result<()> {
        self.apply_cap("cuu1", &[])
    }

    fn delete_line(&mut self) -> Result<()> {
        self.apply_cap("dl", &[])
    }

    fn carriage_return(&mut self) -> Result<()> {
        self.apply_cap("cr", &[])
    }

    fn get_ref<'a>(&'a self) -> &'a T {
        &self.out
    }

    fn get_mut<'a>(&'a mut self) -> &'a mut T {
        &mut self.out
    }

    fn into_inner(self) -> T
        where Self: Sized
    {
        self.out
    }
}

impl<T: Write + Send> TerminfoTerminal<T> {
    /// Create a new TerminfoTerminal with the given TermInfo and Write.
    pub fn new_with_terminfo(out: T, terminfo: TermInfo) -> TerminfoTerminal<T> {
        let nc = if terminfo.strings.contains_key("setaf") &&
                    terminfo.strings.contains_key("setab") {
            terminfo.numbers.get("colors").map_or(0, |&n| n)
        } else {
            0
        };

        TerminfoTerminal {
            out: out,
            ti: terminfo,
            num_colors: nc,
        }
    }

    /// Create a new TerminfoTerminal for the current environment with the given Write.
    ///
    /// Returns `None` when the terminfo cannot be found or parsed.
    pub fn new(out: T) -> Option<TerminfoTerminal<T>> {
        TermInfo::from_env().map(move |ti| TerminfoTerminal::new_with_terminfo(out, ti)).ok()
    }

    fn dim_if_necessary(&self, color: color::Color) -> color::Color {
        if color >= self.num_colors && color >= 8 && color < 16 {
            color - 8
        } else {
            color
        }
    }

    fn apply_cap(&mut self, cmd: &str, params: &[Param]) -> Result<()> {
        match self.ti.strings.get(cmd) {
            Some(cmd) => {
                match expand(&cmd, params, &mut Variables::new()) {
                    Ok(s) => {
                        try!(self.out.write_all(&s));
                        Ok(())
                    }
                    Err(e) => Err(e.into()),
                }
            }
            None => Err(::Error::NotSupported),
        }
    }
}

impl<T: Write> Write for TerminfoTerminal<T> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.out.write(buf)
    }

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


<commit_before>use crate::{
    config, db,
    models::Version,
    schema::{crates, readme_renderings, versions},
    uploaders::Uploader,
};
use anyhow::{anyhow, Context};
use std::{io::Read, path::Path, sync::Arc, thread};

use cargo_registry_markdown::text_to_html;
use chrono::{TimeZone, Utc};
use diesel::{dsl::any, prelude::*};
use flate2::read::GzDecoder;
use reqwest::{blocking::Client, header};
use tar::{self, Archive};

const CACHE_CONTROL_README: &str = "public,max-age=604800";
const USER_AGENT: &str = "crates-admin";

#[derive(clap::Parser, Debug)]
#[clap(
    name = "render-readmes",
    about = "Iterates over every crate versions ever uploaded and (re-)renders their \
        readme using the readme renderer from the cargo_registry crate.",
    after_help = "Warning: this can take a lot of time."
)]
pub struct Opts {
    /// How many versions should be queried and processed at a time.
    #[clap(long, default_value = "25")]
    page_size: usize,

    /// Only rerender readmes that are older than this date.
    #[clap(long)]
    older_than: Option<String>,

    /// Only rerender readmes for the specified crate.
    #[clap(long = "crate")]
    crate_name: Option<String>,
}

pub fn run(opts: Opts) {
    let base_config = Arc::new(config::Base::from_environment());
    let conn = db::connect_now().unwrap();

    let start_time = Utc::now();

    let older_than = if let Some(ref time) = opts.older_than {
        Utc.datetime_from_str(time, "%Y-%m-%d %H:%M:%S")
            .expect("Could not parse --older-than argument as a time")
    } else {
        start_time
    };
    let older_than = older_than.naive_utc();

    println!("Start time:                   {}", start_time);
    println!("Rendering readmes older than: {}", older_than);

    let mut query = versions::table
        .inner_join(crates::table)
        .left_outer_join(readme_renderings::table)
        .filter(
            readme_renderings::rendered_at
                .lt(older_than)
                .or(readme_renderings::version_id.is_null()),
        )
        .select(versions::id)
        .into_boxed();

    if let Some(crate_name) = opts.crate_name {
        println!("Rendering readmes for {}", crate_name);
        query = query.filter(crates::name.eq(crate_name));
    }

    let version_ids: Vec<i32> = query.load(&conn).expect("error loading version ids");

    let total_versions = version_ids.len();
    println!("Rendering {} versions", total_versions);

    let page_size = opts.page_size;

    let total_pages = total_versions / page_size;
    let total_pages = if total_versions % page_size == 0 {
        total_pages
    } else {
        total_pages + 1
    };

    let client = Client::new();

    for (page_num, version_ids_chunk) in version_ids.chunks(page_size).enumerate() {
        println!(
            "= Page {} of {} ==================================",
            page_num + 1,
            total_pages
        );

        let versions: Vec<(Version, String)> = versions::table
            .inner_join(crates::table)
            .filter(versions::id.eq(any(version_ids_chunk)))
            .select((versions::all_columns, crates::name))
            .load(&conn)
            .expect("error loading versions");

        let mut tasks = Vec::with_capacity(page_size as usize);
        for (version, krate_name) in versions {
            Version::record_readme_rendering(version.id, &conn).unwrap_or_else(|_| {
                panic!(
                    "[{}-{}] Couldn't record rendering time",
                    krate_name, version.num
                )
            });
            let client = client.clone();
            let base_config = base_config.clone();
            let handle = thread::spawn(move || {
                println!("[{}-{}] Rendering README...", krate_name, version.num);
                let readme = get_readme(base_config.uploader(), &client, &version, &krate_name);
                if readme.is_none() {
                    return;
                }
                let readme = readme.unwrap();
                let content_length = readme.len() as u64;
                let content = std::io::Cursor::new(readme);
                let readme_path = format!("readmes/{0}/{0}-{1}.html", krate_name, version.num);
                let mut extra_headers = header::HeaderMap::new();
                extra_headers.insert(
                    header::CACHE_CONTROL,
                    header::HeaderValue::from_static(CACHE_CONTROL_README),
                );
                base_config
                    .uploader()
                    .upload(
                        &client,
                        &readme_path,
                        content,
                        content_length,
                        "text/html",
                        extra_headers,
                    )
                    .unwrap_or_else(|_| {
                        panic!(
                            "[{}-{}] Couldn't upload file to S3",
                            krate_name, version.num
                        )
                    });
            });
            tasks.push(handle);
        }
        for handle in tasks {
            if let Err(err) = handle.join() {
                println!("Thread panicked: {:?}", err);
            }
        }
    }
}

/// Renders the readme of an uploaded crate version.
fn get_readme(
    uploader: &Uploader,
    client: &Client,
    version: &Version,
    krate_name: &str,
) -> Option<String> {
    let pkg_name = format!("{}-{}", krate_name, version.num);

    let location = uploader.crate_location(krate_name, &version.num.to_string());

    let location = match uploader {
        Uploader::S3 { .. } => location,
        Uploader::Local => format!("http://localhost:8888/{}", location),
    };

    let mut extra_headers = header::HeaderMap::new();
    extra_headers.insert(
        header::USER_AGENT,
        header::HeaderValue::from_static(USER_AGENT),
    );
    let response = match client.get(&location).headers(extra_headers).send() {
        Ok(r) => r,
        Err(err) => {
            println!("[{}] Unable to fetch crate: {}", pkg_name, err);
            return None;
        }
    };

    if !response.status().is_success() {
        println!(
            "[{}] Failed to get a 200 response: {}",
            pkg_name,
            response.text().unwrap()
        );
        return None;
    }

    let reader = GzDecoder::new(response);
    let archive = Archive::new(reader);
    render_pkg_readme(archive, &pkg_name).ok()
}

fn render_pkg_readme<R: Read>(mut archive: Archive<R>, pkg_name: &str) -> anyhow::Result<String> {
    let mut entries = archive.entries().context("Invalid tar archive entries")?;

    let manifest: Manifest = {
        let path = format!("{}/Cargo.toml", pkg_name);
        let contents = find_file_by_path(&mut entries, Path::new(&path))
            .context("Failed to read Cargo.toml file")?;

        toml::from_str(&contents).context("Failed to parse manifest file")?
    };

    let rendered = {
        let readme_path = manifest
            .package
            .readme
            .clone()
            .unwrap_or_else(|| "README.md".into());
        let path = format!("{}/{}", pkg_name, readme_path);
        let contents = find_file_by_path(&mut entries, Path::new(&path))
            .with_context(|| format!("Failed to read {} file", readme_path))?;

        text_to_html(
            &contents,
            &readme_path,
            manifest.package.repository.as_deref(),
        )
    };
    return Ok(rendered);

    #[derive(Debug, Deserialize)]
    struct Package {
        readme: Option<String>,
        repository: Option<String>,
    }

    #[derive(Debug, Deserialize)]
    struct Manifest {
        package: Package,
    }
}

/// Search an entry by its path in a Tar archive.
fn find_file_by_path<R: Read>(
    entries: &mut tar::Entries<'_, R>,
    path: &Path,
) -> anyhow::Result<String> {
    let mut file = entries
        .filter_map(|entry| entry.ok())
        .find(|file| match file.path() {
            Ok(p) => p == path,
            Err(_) => false,
        })
        .ok_or_else(|| anyhow!("Failed to find tarball entry: {}", path.display()))?;

    let mut contents = String::new();
    file.read_to_string(&mut contents)
        .context("Failed to read file contents")?;

    Ok(contents)
}

#[cfg(test)]
pub mod tests {
    use std::io::Write;
    use tar;

    use super::render_pkg_readme;

    pub fn add_file<W: Write>(pkg: &mut tar::Builder<W>, path: &str, content: &[u8]) {
        let mut header = tar::Header::new_gnu();
        header.set_size(content.len() as u64);
        header.set_cksum();
        pkg.append_data(&mut header, path, content).unwrap();
    }

    #[test]
    fn test_render_pkg_readme() {
        let mut pkg = tar::Builder::new(vec![]);
        add_file(
            &mut pkg,
            "foo-0.0.1/Cargo.toml",
            br#"
[package]
readme = "README.md"
"#,
        );
        add_file(&mut pkg, "foo-0.0.1/README.md", b"readme");
        let serialized_archive = pkg.into_inner().unwrap();
        let result =
            render_pkg_readme(tar::Archive::new(&*serialized_archive), "foo-0.0.1").unwrap();
        assert!(result.contains("readme"))
    }

    #[test]
    fn test_render_pkg_no_readme() {
        let mut pkg = tar::Builder::new(vec![]);
        add_file(
            &mut pkg,
            "foo-0.0.1/Cargo.toml",
            br#"
[package]
"#,
        );
        let serialized_archive = pkg.into_inner().unwrap();
        assert_err!(render_pkg_readme(
            tar::Archive::new(&*serialized_archive),
            "foo-0.0.1"
        ));
    }

    #[test]
    fn test_render_pkg_implicit_readme() {
        let mut pkg = tar::Builder::new(vec![]);
        add_file(
            &mut pkg,
            "foo-0.0.1/Cargo.toml",
            br#"
[package]
"#,
        );
        add_file(&mut pkg, "foo-0.0.1/README.md", b"readme");
        let serialized_archive = pkg.into_inner().unwrap();
        let result =
            render_pkg_readme(tar::Archive::new(&*serialized_archive), "foo-0.0.1").unwrap();
        assert!(result.contains("readme"))
    }

    #[test]
    fn test_render_pkg_readme_w_link() {
        let mut pkg = tar::Builder::new(vec![]);
        add_file(
            &mut pkg,
            "foo-0.0.1/Cargo.toml",
            br#"
[package]
readme = "README.md"
repository = "https://github.com/foo/foo"
"#,
        );
        add_file(
            &mut pkg,
            "foo-0.0.1/README.md",
            b"readme [link](./Other.md)",
        );
        let serialized_archive = pkg.into_inner().unwrap();
        let result =
            render_pkg_readme(tar::Archive::new(&*serialized_archive), "foo-0.0.1").unwrap();
        assert!(result.contains("\"https://github.com/foo/foo/blob/HEAD/./Other.md\""))
    }

    #[test]
    fn test_render_pkg_readme_not_at_root() {
        let mut pkg = tar::Builder::new(vec![]);
        add_file(
            &mut pkg,
            "foo-0.0.1/Cargo.toml",
            br#"
[package]
readme = "docs/README.md"
repository = "https://github.com/foo/foo"
"#,
        );
        add_file(
            &mut pkg,
            "foo-0.0.1/docs/README.md",
            b"docs/readme [link](./Other.md)",
        );
        let serialized_archive = pkg.into_inner().unwrap();
        let result =
            render_pkg_readme(tar::Archive::new(&*serialized_archive), "foo-0.0.1").unwrap();
        assert!(result.contains("docs/readme"));
        assert!(result.contains("\"https://github.com/foo/foo/blob/HEAD/docs/./Other.md\""))
    }
}
<commit_msg>admin::render_readmes: Return `Result` from `get_readme()`<commit_after>use crate::{
    config, db,
    models::Version,
    schema::{crates, readme_renderings, versions},
    uploaders::Uploader,
};
use anyhow::{anyhow, Context};
use std::{io::Read, path::Path, sync::Arc, thread};

use cargo_registry_markdown::text_to_html;
use chrono::{TimeZone, Utc};
use diesel::{dsl::any, prelude::*};
use flate2::read::GzDecoder;
use reqwest::{blocking::Client, header};
use tar::{self, Archive};

const CACHE_CONTROL_README: &str = "public,max-age=604800";
const USER_AGENT: &str = "crates-admin";

#[derive(clap::Parser, Debug)]
#[clap(
    name = "render-readmes",
    about = "Iterates over every crate versions ever uploaded and (re-)renders their \
        readme using the readme renderer from the cargo_registry crate.",
    after_help = "Warning: this can take a lot of time."
)]
pub struct Opts {
    /// How many versions should be queried and processed at a time.
    #[clap(long, default_value = "25")]
    page_size: usize,

    /// Only rerender readmes that are older than this date.
    #[clap(long)]
    older_than: Option<String>,

    /// Only rerender readmes for the specified crate.
    #[clap(long = "crate")]
    crate_name: Option<String>,
}

pub fn run(opts: Opts) {
    let base_config = Arc::new(config::Base::from_environment());
    let conn = db::connect_now().unwrap();

    let start_time = Utc::now();

    let older_than = if let Some(ref time) = opts.older_than {
        Utc.datetime_from_str(time, "%Y-%m-%d %H:%M:%S")
            .expect("Could not parse --older-than argument as a time")
    } else {
        start_time
    };
    let older_than = older_than.naive_utc();

    println!("Start time:                   {}", start_time);
    println!("Rendering readmes older than: {}", older_than);

    let mut query = versions::table
        .inner_join(crates::table)
        .left_outer_join(readme_renderings::table)
        .filter(
            readme_renderings::rendered_at
                .lt(older_than)
                .or(readme_renderings::version_id.is_null()),
        )
        .select(versions::id)
        .into_boxed();

    if let Some(crate_name) = opts.crate_name {
        println!("Rendering readmes for {}", crate_name);
        query = query.filter(crates::name.eq(crate_name));
    }

    let version_ids: Vec<i32> = query.load(&conn).expect("error loading version ids");

    let total_versions = version_ids.len();
    println!("Rendering {} versions", total_versions);

    let page_size = opts.page_size;

    let total_pages = total_versions / page_size;
    let total_pages = if total_versions % page_size == 0 {
        total_pages
    } else {
        total_pages + 1
    };

    let client = Client::new();

    for (page_num, version_ids_chunk) in version_ids.chunks(page_size).enumerate() {
        println!(
            "= Page {} of {} ==================================",
            page_num + 1,
            total_pages
        );

        let versions: Vec<(Version, String)> = versions::table
            .inner_join(crates::table)
            .filter(versions::id.eq(any(version_ids_chunk)))
            .select((versions::all_columns, crates::name))
            .load(&conn)
            .expect("error loading versions");

        let mut tasks = Vec::with_capacity(page_size as usize);
        for (version, krate_name) in versions {
            Version::record_readme_rendering(version.id, &conn).unwrap_or_else(|_| {
                panic!(
                    "[{}-{}] Couldn't record rendering time",
                    krate_name, version.num
                )
            });
            let client = client.clone();
            let base_config = base_config.clone();
            let handle = thread::spawn(move || {
                println!("[{}-{}] Rendering README...", krate_name, version.num);
                let readme = get_readme(base_config.uploader(), &client, &version, &krate_name);
                if readme.is_err() {
                    return;
                }
                let readme = readme.unwrap();
                let content_length = readme.len() as u64;
                let content = std::io::Cursor::new(readme);
                let readme_path = format!("readmes/{0}/{0}-{1}.html", krate_name, version.num);
                let mut extra_headers = header::HeaderMap::new();
                extra_headers.insert(
                    header::CACHE_CONTROL,
                    header::HeaderValue::from_static(CACHE_CONTROL_README),
                );
                base_config
                    .uploader()
                    .upload(
                        &client,
                        &readme_path,
                        content,
                        content_length,
                        "text/html",
                        extra_headers,
                    )
                    .unwrap_or_else(|_| {
                        panic!(
                            "[{}-{}] Couldn't upload file to S3",
                            krate_name, version.num
                        )
                    });
            });
            tasks.push(handle);
        }
        for handle in tasks {
            if let Err(err) = handle.join() {
                println!("Thread panicked: {:?}", err);
            }
        }
    }
}

/// Renders the readme of an uploaded crate version.
fn get_readme(
    uploader: &Uploader,
    client: &Client,
    version: &Version,
    krate_name: &str,
) -> anyhow::Result<String> {
    let pkg_name = format!("{}-{}", krate_name, version.num);

    let location = uploader.crate_location(krate_name, &version.num.to_string());

    let location = match uploader {
        Uploader::S3 { .. } => location,
        Uploader::Local => format!("http://localhost:8888/{}", location),
    };

    let mut extra_headers = header::HeaderMap::new();
    extra_headers.insert(
        header::USER_AGENT,
        header::HeaderValue::from_static(USER_AGENT),
    );
    let request = client.get(&location).headers(extra_headers);
    let response = request.send().context("Failed to fetch crate")?;

    if !response.status().is_success() {
        return Err(anyhow!(
            "Failed to get a 200 response: {}",
            response.text().unwrap()
        ));
    }

    let reader = GzDecoder::new(response);
    let archive = Archive::new(reader);
    render_pkg_readme(archive, &pkg_name)
}

fn render_pkg_readme<R: Read>(mut archive: Archive<R>, pkg_name: &str) -> anyhow::Result<String> {
    let mut entries = archive.entries().context("Invalid tar archive entries")?;

    let manifest: Manifest = {
        let path = format!("{}/Cargo.toml", pkg_name);
        let contents = find_file_by_path(&mut entries, Path::new(&path))
            .context("Failed to read Cargo.toml file")?;

        toml::from_str(&contents).context("Failed to parse manifest file")?
    };

    let rendered = {
        let readme_path = manifest
            .package
            .readme
            .clone()
            .unwrap_or_else(|| "README.md".into());
        let path = format!("{}/{}", pkg_name, readme_path);
        let contents = find_file_by_path(&mut entries, Path::new(&path))
            .with_context(|| format!("Failed to read {} file", readme_path))?;

        text_to_html(
            &contents,
            &readme_path,
            manifest.package.repository.as_deref(),
        )
    };
    return Ok(rendered);

    #[derive(Debug, Deserialize)]
    struct Package {
        readme: Option<String>,
        repository: Option<String>,
    }

    #[derive(Debug, Deserialize)]
    struct Manifest {
        package: Package,
    }
}

/// Search an entry by its path in a Tar archive.
fn find_file_by_path<R: Read>(
    entries: &mut tar::Entries<'_, R>,
    path: &Path,
) -> anyhow::Result<String> {
    let mut file = entries
        .filter_map(|entry| entry.ok())
        .find(|file| match file.path() {
            Ok(p) => p == path,
            Err(_) => false,
        })
        .ok_or_else(|| anyhow!("Failed to find tarball entry: {}", path.display()))?;

    let mut contents = String::new();
    file.read_to_string(&mut contents)
        .context("Failed to read file contents")?;

    Ok(contents)
}

#[cfg(test)]
pub mod tests {
    use std::io::Write;
    use tar;

    use super::render_pkg_readme;

    pub fn add_file<W: Write>(pkg: &mut tar::Builder<W>, path: &str, content: &[u8]) {
        let mut header = tar::Header::new_gnu();
        header.set_size(content.len() as u64);
        header.set_cksum();
        pkg.append_data(&mut header, path, content).unwrap();
    }

    #[test]
    fn test_render_pkg_readme() {
        let mut pkg = tar::Builder::new(vec![]);
        add_file(
            &mut pkg,
            "foo-0.0.1/Cargo.toml",
            br#"
[package]
readme = "README.md"
"#,
        );
        add_file(&mut pkg, "foo-0.0.1/README.md", b"readme");
        let serialized_archive = pkg.into_inner().unwrap();
        let result =
            render_pkg_readme(tar::Archive::new(&*serialized_archive), "foo-0.0.1").unwrap();
        assert!(result.contains("readme"))
    }

    #[test]
    fn test_render_pkg_no_readme() {
        let mut pkg = tar::Builder::new(vec![]);
        add_file(
            &mut pkg,
            "foo-0.0.1/Cargo.toml",
            br#"
[package]
"#,
        );
        let serialized_archive = pkg.into_inner().unwrap();
        assert_err!(render_pkg_readme(
            tar::Archive::new(&*serialized_archive),
            "foo-0.0.1"
        ));
    }

    #[test]
    fn test_render_pkg_implicit_readme() {
        let mut pkg = tar::Builder::new(vec![]);
        add_file(
            &mut pkg,
            "foo-0.0.1/Cargo.toml",
            br#"
[package]
"#,
        );
        add_file(&mut pkg, "foo-0.0.1/README.md", b"readme");
        let serialized_archive = pkg.into_inner().unwrap();
        let result =
            render_pkg_readme(tar::Archive::new(&*serialized_archive), "foo-0.0.1").unwrap();
        assert!(result.contains("readme"))
    }

    #[test]
    fn test_render_pkg_readme_w_link() {
        let mut pkg = tar::Builder::new(vec![]);
        add_file(
            &mut pkg,
            "foo-0.0.1/Cargo.toml",
            br#"
[package]
readme = "README.md"
repository = "https://github.com/foo/foo"
"#,
        );
        add_file(
            &mut pkg,
            "foo-0.0.1/README.md",
            b"readme [link](./Other.md)",
        );
        let serialized_archive = pkg.into_inner().unwrap();
        let result =
            render_pkg_readme(tar::Archive::new(&*serialized_archive), "foo-0.0.1").unwrap();
        assert!(result.contains("\"https://github.com/foo/foo/blob/HEAD/./Other.md\""))
    }

    #[test]
    fn test_render_pkg_readme_not_at_root() {
        let mut pkg = tar::Builder::new(vec![]);
        add_file(
            &mut pkg,
            "foo-0.0.1/Cargo.toml",
            br#"
[package]
readme = "docs/README.md"
repository = "https://github.com/foo/foo"
"#,
        );
        add_file(
            &mut pkg,
            "foo-0.0.1/docs/README.md",
            b"docs/readme [link](./Other.md)",
        );
        let serialized_archive = pkg.into_inner().unwrap();
        let result =
            render_pkg_readme(tar::Archive::new(&*serialized_archive), "foo-0.0.1").unwrap();
        assert!(result.contains("docs/readme"));
        assert!(result.contains("\"https://github.com/foo/foo/blob/HEAD/docs/./Other.md\""))
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2015, Mikhail Vorotilov
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// * Redistributions of source code must retain the above copyright notice, this
//   list of conditions and the following disclaimer.
//
// * Redistributions in binary form must reproduce the above copyright notice,
//   this list of conditions and the following disclaimer in the documentation
//   and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

use super::super::FloatType;
use super::super::Roots;

/// Solves a quadratic equation a2*x^2 + a1*x + a0 = 0.
///
/// In case two roots are present, the first returned root is less than the second one.
///
/// # Examples
///
/// ```
/// use roots::Roots;
/// use roots::find_roots_quadratic;
///
/// let no_roots = find_roots_quadratic(1f32, 0f32, 1f32);
/// // Returns Roots::No([]) as 'x^2 + 1 = 0' has no roots
///
/// let one_root = find_roots_quadratic(1f64, 0f64, 0f64);
/// // Returns Roots::One([0f64]) as 'x^2 = 0' has one root 0
///
/// let two_roots = find_roots_quadratic(1f32, 0f32, -1f32);
/// // Returns Roots::Two([-1f32,1f32]) as 'x^2 - 1 = 0' has roots -1 and 1
/// ```
pub fn find_roots_quadratic<F: FloatType>(a2: F, a1: F, a0: F) -> Roots<F> {
    // Handle non-standard cases
    if a2 == F::zero() {
        // a2 = 0; a1*x+a0=0; solve linear equation
        super::linear::find_roots_linear(a1, a0)
    } else {
        // Rust lacks a simple way to convert an integer constant to generic type F
        let a2x2 = F::two() * a2;
        let discriminant = a1 * a1 - F::four() * a2 * a0;
        if discriminant < F::zero() {
            Roots::No([])
        } else {
            if discriminant == F::zero() {
                Roots::One([-a1 / a2x2])
            } else {
                let sq = discriminant.sqrt();
                Roots::Two([(-a1 - sq) / a2x2, (-a1 + sq) / a2x2])
            }
        }
    }
}

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

    #[test]
    fn test_find_roots_quadratic() {
        assert_eq!(find_roots_quadratic(0f32, 0f32, 0f32), Roots::One([0f32]));
        assert_eq!(find_roots_quadratic(1f32, 0f32, 1f32), Roots::No([]));
        assert_eq!(find_roots_quadratic(1f64, 0f64, -1f64),
                   Roots::Two([-1f64, 1f64]));
    }

    #[test]
    fn test_find_roots_quadratic_small_a2() {
        assert_eq!(find_roots_quadratic(1e-20f32, -1f32, -1e-30f32),
                   Roots::Two([-1e-30f32, 1e20f32]));
        assert_eq!(find_roots_quadratic(-1e-20f32, 1f32, 1e-30f32),
                   Roots::Two([-1e-30f32, 1e20f32]));
        assert_eq!(find_roots_quadratic(1e-20f32, -1f32, 1f32),
                   Roots::Two([1f32, 1e20f32]));
        assert_eq!(find_roots_quadratic(-1e-20f32, 1f32, 1f32),
                   Roots::Two([-1f32, 1e20f32]));
        assert_eq!(find_roots_quadratic(-1e-20f32, 1f32, -1f32),
                   Roots::Two([1f32, 1e20f32]));
    }

    #[test]
    fn test_find_roots_quadratic_big_a1() {
        assert_eq!(find_roots_quadratic(1f32, -1e15f32, -1f32),
                   Roots::Two([-1e-15f32, 1e15f32]));
        assert_eq!(find_roots_quadratic(-1f32, 1e15f32, 1f32),
                   Roots::Two([-1e-15f32, 1e15f32]));
    }
}
<commit_msg>Chose greatest divisors when solving the quadratic equations.<commit_after>// Copyright (c) 2015, Mikhail Vorotilov
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// * Redistributions of source code must retain the above copyright notice, this
//   list of conditions and the following disclaimer.
//
// * Redistributions in binary form must reproduce the above copyright notice,
//   this list of conditions and the following disclaimer in the documentation
//   and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

use super::super::FloatType;
use super::super::Roots;

/// Solves a quadratic equation a2*x^2 + a1*x + a0 = 0.
///
/// In case two roots are present, the first returned root is less than the second one.
///
/// # Examples
///
/// ```
/// use roots::Roots;
/// use roots::find_roots_quadratic;
///
/// let no_roots = find_roots_quadratic(1f32, 0f32, 1f32);
/// // Returns Roots::No([]) as 'x^2 + 1 = 0' has no roots
///
/// let one_root = find_roots_quadratic(1f64, 0f64, 0f64);
/// // Returns Roots::One([0f64]) as 'x^2 = 0' has one root 0
///
/// let two_roots = find_roots_quadratic(1f32, 0f32, -1f32);
/// // Returns Roots::Two([-1f32,1f32]) as 'x^2 - 1 = 0' has roots -1 and 1
/// ```
pub fn find_roots_quadratic<F: FloatType>(a2: F, a1: F, a0: F) -> Roots<F> {
    // Handle non-standard cases
    if a2 == F::zero() {
        // a2 = 0; a1*x+a0=0; solve linear equation
        super::linear::find_roots_linear(a1, a0)
    } else {
        // Rust lacks a simple way to convert an integer constant to generic type F
        let discriminant = a1 * a1 - F::four() * a2 * a0;
        if discriminant < F::zero() {
            Roots::No([])
        } else {
            let a2x2 = F::two() * a2;
            if discriminant == F::zero() {
                Roots::One([-a1 / a2x2])
            } else {
                // To improve precision, do not use the smallest divisor.
                // See https://people.csail.mit.edu/bkph/articles/Quadratics.pdf
                let sq = discriminant.sqrt();

                let (same_sign, diff_sign) = if a1 < F::zero() {
                    (-a1 + sq, -a1 - sq)
                } else {
                    (-a1 - sq, -a1 + sq)
                };

                let (x1, x2) = if same_sign.abs() > a2x2.abs() {
                    let a0x2 = F::two() * a0;
                    if diff_sign.abs() > a2x2.abs() {
                        // 2*a2 is the smallest divisor, do not use it
                        (a0x2 / same_sign, a0x2 / diff_sign)
                    } else {
                        // diff_sign is the smallest divisor, do not use it
                        (a0x2 / same_sign, same_sign / a2x2)
                    }
                } else {
                    // 2*a2 is the greatest divisor, use it
                    (diff_sign / a2x2, same_sign / a2x2)
                };

                // Order roots
                if x1 < x2 {
                    Roots::Two([x1, x2])
                } else {
                    Roots::Two([x2, x1])
                }
            }
        }
    }
}

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

    #[test]
    fn test_find_roots_quadratic() {
        assert_eq!(find_roots_quadratic(0f32, 0f32, 0f32), Roots::One([0f32]));
        assert_eq!(find_roots_quadratic(1f32, 0f32, 1f32), Roots::No([]));
        assert_eq!(find_roots_quadratic(1f64, 0f64, -1f64),
                   Roots::Two([-1f64, 1f64]));
    }

    #[test]
    fn test_find_roots_quadratic_small_a2() {
        assert_eq!(find_roots_quadratic(1e-20f32, -1f32, -1e-30f32),
                   Roots::Two([-1e-30f32, 1e20f32]));
        assert_eq!(find_roots_quadratic(-1e-20f32, 1f32, 1e-30f32),
                   Roots::Two([-1e-30f32, 1e20f32]));
        assert_eq!(find_roots_quadratic(1e-20f32, -1f32, 1f32),
                   Roots::Two([1f32, 1e20f32]));
        assert_eq!(find_roots_quadratic(-1e-20f32, 1f32, 1f32),
                   Roots::Two([-1f32, 1e20f32]));
        assert_eq!(find_roots_quadratic(-1e-20f32, 1f32, -1f32),
                   Roots::Two([1f32, 1e20f32]));
    }

    #[test]
    fn test_find_roots_quadratic_big_a1() {
        assert_eq!(find_roots_quadratic(1f32, -1e15f32, -1f32),
                   Roots::Two([-1e-15f32, 1e15f32]));
        assert_eq!(find_roots_quadratic(-1f32, 1e15f32, 1f32),
                   Roots::Two([-1e-15f32, 1e15f32]));
    }
}
<|endoftext|>