<commit_before>use token::{Token, TokenType};
use ast::*;
use pos::Pos;
use error::Error;
use types::Type;

pub struct Parser {
    tokens: Vec<Token>,
    index: usize,
}

impl Parser {
    pub fn new(tokens: Vec<Token>) -> Self {
        Parser {
            tokens: tokens,
            index: 0,
        }
    }

    fn peek(&self, t: TokenType) -> bool {
        let token = &self.tokens[self.index];
        token.typ == t
    }

    fn curr_token(&self) -> Token {
        self.tokens[self.index].clone()
    }

    fn token_pos(&self) -> Pos {
        self.tokens[self.index].pos
    }

    fn eat(&mut self, t: TokenType) -> Result<(), Error> {
        let b = self.peek(t);
        if b {
            self.index += 1;
            Ok(())
        } else {
            Err(Error::UnexpectedToken(self.curr_token(), vec![t]))
        }
    }

    fn eat_lexeme(&mut self, t: TokenType) -> Result<String, Error> {
        let token = &self.tokens[self.index];
        if token.typ == t {
            match token.lexeme {
                Some(ref lexeme) => {
                    self.index += 1;
                    Ok(lexeme.clone())
                }
                None => Err(Error::UnexpectedToken(self.curr_token(), vec![t]))
            }
        } else {
            Err(Error::UnexpectedToken(self.curr_token(), vec![t]))
        }
    }

    fn eat_type(&mut self) -> Result<Type, Error> {
        if self.peek(TokenType::TypeInt) {
            self.index += 1;
            Ok(Type::Int)
        } else if self.peek(TokenType::TypeFloat) {
            self.index += 1;
            Ok(Type::Float)
        } else {
            Err(Error::UnexpectedToken(self.curr_token(), vec![TokenType::TypeInt, TokenType::TypeFloat]))
        }
    }

    pub fn parse_program(&mut self) -> Result<Program, Error> {
        let decls = try!(self.parse_decls());
        let stmts = try!(self.parse_stmts());
        try!(self.eat(TokenType::Eof));

        Ok(Program {
            decls: decls,
            stmts: stmts,
        })
    }

    fn parse_decls(&mut self) -> Result<Vec<Decl>, Error> {
        let mut decls: Vec<Decl> = Vec::new();
        while self.peek(TokenType::Var) || self.peek(TokenType::Function) {
            let decl = try!(self.parse_decl());
            decls.push(decl);
        }
        Ok(decls)
    }

    fn parse_decl(&mut self) -> Result<Decl, Error> {
        if self.peek(TokenType::Var) {
            let vd = try!(self.parse_var_decl());
            Ok(Decl::Var(vd))
        } else if self.peek(TokenType::Function) {
            let fd = try!(self.parse_func_decl());
            Ok(Decl::Fun(fd))
        } else {
            Err(Error::UnexpectedToken(
                self.curr_token(),
                vec![TokenType::Var, TokenType::Function]))
        }
    }

    fn parse_var_decls(&mut self) -> Result<Vec<VarDecl>, Error> {
        let mut vds = Vec::new();
        while self.peek(TokenType::Var) {
            let vd = try!(self.parse_var_decl());
            vds.push(vd);
        }
        Ok(vds)
    }

    fn parse_var_decl(&mut self) -> Result<VarDecl, Error> {
        let pos = self.token_pos();
        try!(self.eat(TokenType::Var));
        let id = try!(self.eat_lexeme(TokenType::Id));
        try!(self.eat(TokenType::Colon));
        let ty = try!(self.eat_type());
        try!(self.eat(TokenType::Semicolon));
        Ok(VarDecl { pos: pos, id: id, ty: ty })
    }

    fn parse_func_decl(&mut self) -> Result<FunDecl, Error> {
        let pos = self.token_pos();
        try!(self.eat(TokenType::Function));
        let id = try!(self.eat_lexeme(TokenType::Id));
        try!(self.eat(TokenType::LParen));
        let params = try!(self.parse_params());
        try!(self.eat(TokenType::RParen));
        try!(self.eat(TokenType::Colon));
        let ty = try!(self.eat_type());
        let decls = try!(self.parse_var_decls());
        let stmts = try!(self.parse_stmts());
        try!(self.eat(TokenType::End));

        Ok(FunDecl {
            pos: pos,
            id: id,
            params: params,
            ty: ty,
            decls: decls,
            stmts: stmts,
        })
    }

    fn parse_params(&mut self) -> Result<Vec<(String, Type)>, Error> {
        let mut params = Vec::new();
        while self.peek(TokenType::Id) {
            let id = try!(self.eat_lexeme(TokenType::Id));
            try!(self.eat(TokenType::Colon));
            let ty = try!(self.eat_type());
            params.push((id, ty));

            if self.peek(TokenType::RParen) {
                break;
            } else if self.peek(TokenType::Comma) {
                try!(self.eat(TokenType::Comma));
            } else {
                return Err(Error::UnexpectedToken(
                    self.curr_token(),
                    vec![TokenType::Comma, TokenType::RParen]));
            }
        }
        Ok(params)
    }

    fn parse_stmts(&mut self) -> Result<Vec<Stmt>, Error> {
        let mut stmts: Vec<Stmt> = Vec::new();
        while self.is_stmt_start() {
            let stmt = try!(self.parse_stmt());
            stmts.push(stmt);
        }
        Ok(stmts)
    }

    fn parse_stmt(&mut self) -> Result<Stmt, Error> {
        if self.peek(TokenType::Read) {
            self.parse_read()
        } else if self.peek(TokenType::Print) {
            self.parse_print()
        } else if self.peek(TokenType::Id) {
            self.parse_assign()
        } else if self.peek(TokenType::If) {
            self.parse_if()
        } else if self.peek(TokenType::While) {
            self.parse_while()
        } else if self.peek(TokenType::Return) {
            self.parse_return()
        } else {
            Err(Error::UnexpectedToken(
                self.curr_token(),
                vec![TokenType::Read, TokenType::Print,
                     TokenType::Id, TokenType::If, TokenType::While]))
        }
    }

    fn parse_read(&mut self) -> Result<Stmt, Error> {
        let pos = self.token_pos();
        try!(self.eat(TokenType::Read));
        let id = try!(self.eat_lexeme(TokenType::Id));
        try!(self.eat(TokenType::Semicolon));
        Ok(Stmt::Read(StmtRead { pos: pos, id: id }))
    }

    fn parse_print(&mut self) -> Result<Stmt, Error> {
        let pos = self.token_pos();
        try!(self.eat(TokenType::Print));
        let e = try!(self.parse_expr());
        try!(self.eat(TokenType::Semicolon));
        Ok(Stmt::Print(StmtPrint { pos: pos, expr: e }))
    }

    fn parse_assign(&mut self) -> Result<Stmt, Error> {
        let pos = self.token_pos();
        let id = try!(self.eat_lexeme(TokenType::Id));
        try!(self.eat(TokenType::Equal));
        let e = try!(self.parse_expr());
        try!(self.eat(TokenType::Semicolon));
        Ok(Stmt::Assign(StmtAssign { pos: pos, id: id, expr: e }))
    }

    fn parse_if(&mut self) -> Result<Stmt, Error> {
        let pos = self.token_pos();
        try!(self.eat(TokenType::If));
        let e = try!(self.parse_expr());
        try!(self.eat(TokenType::Then));
        let then_stmts = try!(self.parse_stmts());
        try!(self.eat(TokenType::Else));
        let else_stmts = try!(self.parse_stmts());
        try!(self.eat(TokenType::End));
        Ok(Stmt::If(StmtIf {
            pos: pos,
            expr: e,
            then_stmts: then_stmts,
            else_stmts: else_stmts,
        }))
    }

    fn parse_while(&mut self) -> Result<Stmt, Error> {
        let pos = self.token_pos();
        try!(self.eat(TokenType::While));
        let e = try!(self.parse_expr());
        try!(self.eat(TokenType::Do));
        let stmts = try!(self.parse_stmts());
        try!(self.eat(TokenType::Done));
        Ok(Stmt::While(StmtWhile {
            pos: pos,
            expr: e,
            stmts: stmts,
        }))
    }

    fn parse_return(&mut self) -> Result<Stmt, Error> {
        let pos = self.token_pos();
        try!(self.eat(TokenType::Return));
        let expr = try!(self.parse_expr());
        try!(self.eat(TokenType::Semicolon));
        Ok(Stmt::Return(StmtReturn {
            pos: pos,
            expr: expr,
        }))
    }

    fn parse_expr(&mut self) -> Result<Expr, Error> {
        let pos = self.token_pos();
        let mut term = try!(self.parse_term());
        while self.next_is_add() {
            if self.peek(TokenType::Plus) {
                try!(self.eat(TokenType::Plus));
                let t2 = try!(self.parse_term());
                term = Expr::Binop(ExprBinop {
                    pos: pos,
                    op: Binop::Add,
                    expr1: Box::new(term),
                    expr2: Box::new(t2)
                });
            } else if self.peek(TokenType::Minus) {
                try!(self.eat(TokenType::Minus));
                let t2 = try!(self.parse_term());
                term = Expr::Binop(ExprBinop {
                    pos: pos,
                    op: Binop::Sub,
                    expr1: Box::new(term),
                    expr2: Box::new(t2)
                });
            } else {
                return Err(Error::UnexpectedToken(
                    self.curr_token(),
                    vec![TokenType::Plus, TokenType::Minus]));
            }
        }
        Ok(term)
    }

    fn parse_term(&mut self) -> Result<Expr, Error> {
        let pos = self.token_pos();
        let mut fact = try!(self.parse_factor());
        while self.next_is_mul() {
            if self.peek(TokenType::Star) {
                try!(self.eat(TokenType::Star));
                let f2 = try!(self.parse_factor());
                fact = Expr::Binop(ExprBinop {
                    pos: pos,
                    op: Binop::Mul,
                    expr1: Box::new(fact),
                    expr2: Box::new(f2),
                });
            } else if self.peek(TokenType::Slash) {
                try!(self.eat(TokenType::Slash));
                let f2 = try!(self.parse_factor());
                fact = Expr::Binop(ExprBinop {
                    pos: pos,
                    op: Binop::Div,
                    expr1: Box::new(fact),
                    expr2: Box::new(f2),
                });
            } else {
                return Err(Error::UnexpectedToken(
                    self.curr_token(),
                    vec![TokenType::Star, TokenType::Slash]));
            }
        }
        Ok(fact)
    }

    fn parse_factor(&mut self) -> Result<Expr, Error> {
        let pos = self.token_pos();
        if self.peek(TokenType::Int) {
            self.parse_int()
        } else if self.peek(TokenType::Float) {
            self.parse_float()
        } else if self.peek(TokenType::Id) {
            self.parse_id()
        } else if self.peek(TokenType::LParen) {
            try!(self.eat(TokenType::LParen));
            let e = try!(self.parse_expr());
            try!(self.eat(TokenType::RParen));
            Ok(e)
        } else if self.peek(TokenType::Minus) {
            try!(self.eat(TokenType::Minus));
            let e = try!(self.parse_expr());
            Ok(Expr::Negate(ExprNegate { pos: pos, expr: Box::new(e) }))
        } else {
            Err(Error::UnexpectedToken(
                self.curr_token(),
                vec![TokenType::Int, TokenType::Float, TokenType::Id,
                     TokenType::Minus, TokenType::LParen]))
        }
    }

    fn parse_int(&mut self) -> Result<Expr, Error> {
        let pos = self.token_pos();
        let lexeme = try!(self.eat_lexeme(TokenType::Int));
        match lexeme.parse::<i64>() {
            Ok(n) => Ok(Expr::Int (ExprInt { pos: pos, value: n })),
            Err(_) => Err(Error::InvalidIntLiteral(pos, lexeme))
        }
    }

    fn parse_float(&mut self) -> Result<Expr, Error> {
        let pos = self.token_pos();
        let lexeme = try!(self.eat_lexeme(TokenType::Float));
        match lexeme.parse::<f64>() {
            Ok(n) => Ok(Expr::Float(ExprFloat { pos: pos, value: Float(n) })),
            Err(_) => Err(Error::InvalidFloatLiteral(pos, lexeme))
        }
    }

    fn parse_id(&mut self) -> Result<Expr, Error> {
        let pos = self.token_pos();
        let lexeme = try!(self.eat_lexeme(TokenType::Id));
        Ok(Expr::Id(ExprId { pos: pos, id: lexeme }))
    }

    fn is_stmt_start(&self) -> bool {
        self.peek(TokenType::Id) ||
            self.peek(TokenType::If) ||
            self.peek(TokenType::While) ||
            self.peek(TokenType::Read) ||
            self.peek(TokenType::Print) ||
            self.peek(TokenType::Return)
    }

    fn next_is_add(&self) -> bool {
        self.peek(TokenType::Plus) || self.peek(TokenType::Minus)
    }

    fn next_is_mul(&self) -> bool {
        self.peek(TokenType::Star) || self.peek(TokenType::Slash)
    }
}
<commit_msg>Modify handling of param list<commit_after>use token::{Token, TokenType};
use ast::*;
use pos::Pos;
use error::Error;
use types::Type;

pub struct Parser {
    tokens: Vec<Token>,
    index: usize,
}

impl Parser {
    pub fn new(tokens: Vec<Token>) -> Self {
        Parser {
            tokens: tokens,
            index: 0,
        }
    }

    fn peek(&self, t: TokenType) -> bool {
        let token = &self.tokens[self.index];
        token.typ == t
    }

    fn curr_token(&self) -> Token {
        self.tokens[self.index].clone()
    }

    fn token_pos(&self) -> Pos {
        self.tokens[self.index].pos
    }

    fn eat(&mut self, t: TokenType) -> Result<(), Error> {
        let b = self.peek(t);
        if b {
            self.index += 1;
            Ok(())
        } else {
            Err(Error::UnexpectedToken(self.curr_token(), vec![t]))
        }
    }

    fn eat_lexeme(&mut self, t: TokenType) -> Result<String, Error> {
        let token = &self.tokens[self.index];
        if token.typ == t {
            match token.lexeme {
                Some(ref lexeme) => {
                    self.index += 1;
                    Ok(lexeme.clone())
                }
                None => Err(Error::UnexpectedToken(self.curr_token(), vec![t]))
            }
        } else {
            Err(Error::UnexpectedToken(self.curr_token(), vec![t]))
        }
    }

    fn eat_type(&mut self) -> Result<Type, Error> {
        if self.peek(TokenType::TypeInt) {
            self.index += 1;
            Ok(Type::Int)
        } else if self.peek(TokenType::TypeFloat) {
            self.index += 1;
            Ok(Type::Float)
        } else {
            Err(Error::UnexpectedToken(self.curr_token(), vec![TokenType::TypeInt, TokenType::TypeFloat]))
        }
    }

    pub fn parse_program(&mut self) -> Result<Program, Error> {
        let decls = try!(self.parse_decls());
        let stmts = try!(self.parse_stmts());
        try!(self.eat(TokenType::Eof));

        Ok(Program {
            decls: decls,
            stmts: stmts,
        })
    }

    fn parse_decls(&mut self) -> Result<Vec<Decl>, Error> {
        let mut decls: Vec<Decl> = Vec::new();
        while self.peek(TokenType::Var) || self.peek(TokenType::Function) {
            let decl = try!(self.parse_decl());
            decls.push(decl);
        }
        Ok(decls)
    }

    fn parse_decl(&mut self) -> Result<Decl, Error> {
        if self.peek(TokenType::Var) {
            let vd = try!(self.parse_var_decl());
            Ok(Decl::Var(vd))
        } else if self.peek(TokenType::Function) {
            let fd = try!(self.parse_func_decl());
            Ok(Decl::Fun(fd))
        } else {
            Err(Error::UnexpectedToken(
                self.curr_token(),
                vec![TokenType::Var, TokenType::Function]))
        }
    }

    fn parse_var_decls(&mut self) -> Result<Vec<VarDecl>, Error> {
        let mut vds = Vec::new();
        while self.peek(TokenType::Var) {
            let vd = try!(self.parse_var_decl());
            vds.push(vd);
        }
        Ok(vds)
    }

    fn parse_var_decl(&mut self) -> Result<VarDecl, Error> {
        let pos = self.token_pos();
        try!(self.eat(TokenType::Var));
        let id = try!(self.eat_lexeme(TokenType::Id));
        try!(self.eat(TokenType::Colon));
        let ty = try!(self.eat_type());
        try!(self.eat(TokenType::Semicolon));
        Ok(VarDecl { pos: pos, id: id, ty: ty })
    }

    fn parse_func_decl(&mut self) -> Result<FunDecl, Error> {
        let pos = self.token_pos();
        try!(self.eat(TokenType::Function));
        let id = try!(self.eat_lexeme(TokenType::Id));
        try!(self.eat(TokenType::LParen));
        let params = try!(self.parse_params());
        try!(self.eat(TokenType::RParen));
        try!(self.eat(TokenType::Colon));
        let ty = try!(self.eat_type());
        let decls = try!(self.parse_var_decls());
        let stmts = try!(self.parse_stmts());
        try!(self.eat(TokenType::End));

        Ok(FunDecl {
            pos: pos,
            id: id,
            params: params,
            ty: ty,
            decls: decls,
            stmts: stmts,
        })
    }

    fn parse_params(&mut self) -> Result<Vec<(String, Type)>, Error> {
        let mut params = Vec::new();
        while self.peek(TokenType::Id) {
            let id = try!(self.eat_lexeme(TokenType::Id));
            try!(self.eat(TokenType::Colon));
            let ty = try!(self.eat_type());
            params.push((id, ty));

            if self.peek(TokenType::Comma) {
                try!(self.eat(TokenType::Comma));
            } else {
                break;
            }
        }
        Ok(params)
    }

    fn parse_stmts(&mut self) -> Result<Vec<Stmt>, Error> {
        let mut stmts: Vec<Stmt> = Vec::new();
        while self.is_stmt_start() {
            let stmt = try!(self.parse_stmt());
            stmts.push(stmt);
        }
        Ok(stmts)
    }

    fn parse_stmt(&mut self) -> Result<Stmt, Error> {
        if self.peek(TokenType::Read) {
            self.parse_read()
        } else if self.peek(TokenType::Print) {
            self.parse_print()
        } else if self.peek(TokenType::Id) {
            self.parse_assign()
        } else if self.peek(TokenType::If) {
            self.parse_if()
        } else if self.peek(TokenType::While) {
            self.parse_while()
        } else if self.peek(TokenType::Return) {
            self.parse_return()
        } else {
            Err(Error::UnexpectedToken(
                self.curr_token(),
                vec![TokenType::Read, TokenType::Print,
                     TokenType::Id, TokenType::If, TokenType::While]))
        }
    }

    fn parse_read(&mut self) -> Result<Stmt, Error> {
        let pos = self.token_pos();
        try!(self.eat(TokenType::Read));
        let id = try!(self.eat_lexeme(TokenType::Id));
        try!(self.eat(TokenType::Semicolon));
        Ok(Stmt::Read(StmtRead { pos: pos, id: id }))
    }

    fn parse_print(&mut self) -> Result<Stmt, Error> {
        let pos = self.token_pos();
        try!(self.eat(TokenType::Print));
        let e = try!(self.parse_expr());
        try!(self.eat(TokenType::Semicolon));
        Ok(Stmt::Print(StmtPrint { pos: pos, expr: e }))
    }

    fn parse_assign(&mut self) -> Result<Stmt, Error> {
        let pos = self.token_pos();
        let id = try!(self.eat_lexeme(TokenType::Id));
        try!(self.eat(TokenType::Equal));
        let e = try!(self.parse_expr());
        try!(self.eat(TokenType::Semicolon));
        Ok(Stmt::Assign(StmtAssign { pos: pos, id: id, expr: e }))
    }

    fn parse_if(&mut self) -> Result<Stmt, Error> {
        let pos = self.token_pos();
        try!(self.eat(TokenType::If));
        let e = try!(self.parse_expr());
        try!(self.eat(TokenType::Then));
        let then_stmts = try!(self.parse_stmts());
        try!(self.eat(TokenType::Else));
        let else_stmts = try!(self.parse_stmts());
        try!(self.eat(TokenType::End));
        Ok(Stmt::If(StmtIf {
            pos: pos,
            expr: e,
            then_stmts: then_stmts,
            else_stmts: else_stmts,
        }))
    }

    fn parse_while(&mut self) -> Result<Stmt, Error> {
        let pos = self.token_pos();
        try!(self.eat(TokenType::While));
        let e = try!(self.parse_expr());
        try!(self.eat(TokenType::Do));
        let stmts = try!(self.parse_stmts());
        try!(self.eat(TokenType::Done));
        Ok(Stmt::While(StmtWhile {
            pos: pos,
            expr: e,
            stmts: stmts,
        }))
    }

    fn parse_return(&mut self) -> Result<Stmt, Error> {
        let pos = self.token_pos();
        try!(self.eat(TokenType::Return));
        let expr = try!(self.parse_expr());
        try!(self.eat(TokenType::Semicolon));
        Ok(Stmt::Return(StmtReturn {
            pos: pos,
            expr: expr,
        }))
    }

    fn parse_expr(&mut self) -> Result<Expr, Error> {
        let pos = self.token_pos();
        let mut term = try!(self.parse_term());
        while self.next_is_add() {
            if self.peek(TokenType::Plus) {
                try!(self.eat(TokenType::Plus));
                let t2 = try!(self.parse_term());
                term = Expr::Binop(ExprBinop {
                    pos: pos,
                    op: Binop::Add,
                    expr1: Box::new(term),
                    expr2: Box::new(t2)
                });
            } else if self.peek(TokenType::Minus) {
                try!(self.eat(TokenType::Minus));
                let t2 = try!(self.parse_term());
                term = Expr::Binop(ExprBinop {
                    pos: pos,
                    op: Binop::Sub,
                    expr1: Box::new(term),
                    expr2: Box::new(t2)
                });
            } else {
                return Err(Error::UnexpectedToken(
                    self.curr_token(),
                    vec![TokenType::Plus, TokenType::Minus]));
            }
        }
        Ok(term)
    }

    fn parse_term(&mut self) -> Result<Expr, Error> {
        let pos = self.token_pos();
        let mut fact = try!(self.parse_factor());
        while self.next_is_mul() {
            if self.peek(TokenType::Star) {
                try!(self.eat(TokenType::Star));
                let f2 = try!(self.parse_factor());
                fact = Expr::Binop(ExprBinop {
                    pos: pos,
                    op: Binop::Mul,
                    expr1: Box::new(fact),
                    expr2: Box::new(f2),
                });
            } else if self.peek(TokenType::Slash) {
                try!(self.eat(TokenType::Slash));
                let f2 = try!(self.parse_factor());
                fact = Expr::Binop(ExprBinop {
                    pos: pos,
                    op: Binop::Div,
                    expr1: Box::new(fact),
                    expr2: Box::new(f2),
                });
            } else {
                return Err(Error::UnexpectedToken(
                    self.curr_token(),
                    vec![TokenType::Star, TokenType::Slash]));
            }
        }
        Ok(fact)
    }

    fn parse_factor(&mut self) -> Result<Expr, Error> {
        let pos = self.token_pos();
        if self.peek(TokenType::Int) {
            self.parse_int()
        } else if self.peek(TokenType::Float) {
            self.parse_float()
        } else if self.peek(TokenType::Id) {
            self.parse_id()
        } else if self.peek(TokenType::LParen) {
            try!(self.eat(TokenType::LParen));
            let e = try!(self.parse_expr());
            try!(self.eat(TokenType::RParen));
            Ok(e)
        } else if self.peek(TokenType::Minus) {
            try!(self.eat(TokenType::Minus));
            let e = try!(self.parse_expr());
            Ok(Expr::Negate(ExprNegate { pos: pos, expr: Box::new(e) }))
        } else {
            Err(Error::UnexpectedToken(
                self.curr_token(),
                vec![TokenType::Int, TokenType::Float, TokenType::Id,
                     TokenType::Minus, TokenType::LParen]))
        }
    }

    fn parse_int(&mut self) -> Result<Expr, Error> {
        let pos = self.token_pos();
        let lexeme = try!(self.eat_lexeme(TokenType::Int));
        match lexeme.parse::<i64>() {
            Ok(n) => Ok(Expr::Int (ExprInt { pos: pos, value: n })),
            Err(_) => Err(Error::InvalidIntLiteral(pos, lexeme))
        }
    }

    fn parse_float(&mut self) -> Result<Expr, Error> {
        let pos = self.token_pos();
        let lexeme = try!(self.eat_lexeme(TokenType::Float));
        match lexeme.parse::<f64>() {
            Ok(n) => Ok(Expr::Float(ExprFloat { pos: pos, value: Float(n) })),
            Err(_) => Err(Error::InvalidFloatLiteral(pos, lexeme))
        }
    }

    fn parse_id(&mut self) -> Result<Expr, Error> {
        let pos = self.token_pos();
        let lexeme = try!(self.eat_lexeme(TokenType::Id));
        Ok(Expr::Id(ExprId { pos: pos, id: lexeme }))
    }

    fn is_stmt_start(&self) -> bool {
        self.peek(TokenType::Id) ||
            self.peek(TokenType::If) ||
            self.peek(TokenType::While) ||
            self.peek(TokenType::Read) ||
            self.peek(TokenType::Print) ||
            self.peek(TokenType::Return)
    }

    fn next_is_add(&self) -> bool {
        self.peek(TokenType::Plus) || self.peek(TokenType::Minus)
    }

    fn next_is_mul(&self) -> bool {
        self.peek(TokenType::Star) || self.peek(TokenType::Slash)
    }
}
<|endoftext|>


<commit_before>use xplane_plugin::*;
use xpgl::gl;
use xplm;

use xplm::ui::*;
use xplm::graphics::*;
use xplm::graphics::window::*;

use mapcore::map::Map;
use mapcore::layer::Layer;
use mapcore::{LatLon, LatLonRect};
use mapcore::projection::Projection;
use mapcore::equirectangular::EquirectangularProjection;

use std::rc::Rc;
use std::cell::RefCell;
use std::mem;

/// The graphics state used for 2D drawing
const GRAPHICS_STATE_2D : GraphicsState = GraphicsState {
    fog: false,
    lighting: false,
    alpha_testing: false,
    alpha_blending: true,
    depth_testing: false,
    depth_writing: false,
    textures: 0,
};

pub struct TerrainDisplayPlugin {
    /// The window that displays graphics
    window: Option<Rc<RefCell<Window>>>,
}

impl Plugin for TerrainDisplayPlugin {
    fn start() -> Option<Self> {
        // Load OpenGL symbols
        gl::load_with(|name| {
            unsafe { mem::transmute(xplm::find_symbol(name)) }
        });

        Some(TerrainDisplayPlugin {
            window: None,
        })
    }
    fn enable(&mut self) {
        self.window = Some(Window::new(&Rect { left: 100, top: 1000, right: 400, bottom: 600 }));
        let mut map = Map::new(EquirectangularProjection);
        map.add_layer(TestLayer);
        {
            let local_draw_window = move |window: &mut Window| {

                let rect = window.get_geometry();

                set_state(&GRAPHICS_STATE_2D);
                unsafe {
                    gl::Color3f(0.0, 0.5, 0.0);
                    gl::Begin(gl::QUADS);
                    gl::Vertex2i(rect.left, rect.bottom);
                    gl::Vertex2i(rect.left, rect.top);
                    gl::Vertex2i(rect.right, rect.top);
                    gl::Vertex2i(rect.right, rect.bottom);
                    gl::End();
                }
                map.draw(rect.left, rect.top, (rect.right - rect.left), (rect.bottom - rect.top));
            };

            let mut window = self.window.as_mut().unwrap().borrow_mut();
            window.set_draw_callback(local_draw_window);
            window.set_visible(true);
        }
    }
    fn disable(&mut self) {
        self.window = None;
    }

    fn info<'a, 'b, 'c>(&self) -> PluginInfo<'a, 'b, 'c> {
        PluginInfo {
            name: "Terrain Display",
            signature: "org.samcrow.dev.terrain_display",
            description: "Draws a display of the terrain around the aircraft",
        }
    }

    fn stop(&mut self) {

    }
}

struct TestLayer;

impl Layer for TestLayer {
    fn draw(&self, projection: &Projection, x: i32, y: i32, width: i32, height: i32) {
        let height = height / 2;
        set_state(&GRAPHICS_STATE_2D);
        unsafe {
            gl::Color3f(0.8, 0.3, 0.0);
            gl::Begin(gl::QUADS);
            gl::Vertex2i(x, y + height);
            gl::Vertex2i(x , y );
            gl::Vertex2i(x + width, y );
            gl::Vertex2i(x + width, y + height);
            gl::End();
        }
    }

    fn bounds(&self) -> Option<LatLonRect> {
        None
    }
}
<commit_msg>Fixes<commit_after>use xplane_plugin::*;
use xpgl::gl;
use xplm;

use xplm::ui::*;
use xplm::graphics::*;
use xplm::graphics::window::*;

use mapcore::map::Map;
use mapcore::layer::Layer;
use mapcore::{Latitude, Longitude, LatLon, LatLonRect, Polygon};
use mapcore::projection::Projection;
use mapcore::equirectangular::EquirectangularProjection;

use std::rc::Rc;
use std::cell::RefCell;
use std::mem;

/// The graphics state used for 2D drawing
const GRAPHICS_STATE_2D : GraphicsState = GraphicsState {
    fog: false,
    lighting: false,
    alpha_testing: false,
    alpha_blending: true,
    depth_testing: false,
    depth_writing: false,
    textures: 0,
};

pub struct TerrainDisplayPlugin {
    /// The window that displays graphics
    window: Option<Rc<RefCell<Window>>>,
}

impl Plugin for TerrainDisplayPlugin {
    fn start() -> Option<Self> {
        // Load OpenGL symbols
        gl::load_with(|name| {
            unsafe { mem::transmute(xplm::find_symbol(name)) }
        });

        Some(TerrainDisplayPlugin {
            window: None,
        })
    }
    fn enable(&mut self) {
        self.window = Some(Window::new(&Rect { left: 100, top: 1000, right: 400, bottom: 600 }));
        let mut map = Map::new(EquirectangularProjection);
        map.add_layer(TestLayer);
        {
            let local_draw_window = move |window: &mut Window| {

                let rect = window.get_geometry();

                set_state(&GRAPHICS_STATE_2D);
                unsafe {
                    gl::Color3f(0.0, 0.5, 0.0);
                    gl::Begin(gl::QUADS);
                    gl::Vertex2i(rect.left, rect.bottom);
                    gl::Vertex2i(rect.left, rect.top);
                    gl::Vertex2i(rect.right, rect.top);
                    gl::Vertex2i(rect.right, rect.bottom);
                    gl::End();
                }
                map.draw(rect.left, rect.bottom, (rect.right - rect.left), (rect.top - rect.bottom));
            };

            let mut window = self.window.as_mut().unwrap().borrow_mut();
            window.set_draw_callback(local_draw_window);
            window.set_visible(true);
        }
    }
    fn disable(&mut self) {
        self.window = None;
    }

    fn info<'a, 'b, 'c>(&self) -> PluginInfo<'a, 'b, 'c> {
        PluginInfo {
            name: "Terrain Display",
            signature: "org.samcrow.dev.terrain_display",
            description: "Draws a display of the terrain around the aircraft",
        }
    }

    fn stop(&mut self) {

    }
}

struct TestLayer;

impl Layer for TestLayer {
    fn draw(&self, projection: &Projection, x: i32, y: i32, width: i32, height: i32) {
        set_state(&GRAPHICS_STATE_2D);
        unsafe {
            gl::Color3f(0.8, 0.3, 0.0);
            gl::Begin(gl::QUADS);
            gl::Vertex2i(x, y);
            gl::Vertex2i(x , y + height);
            gl::Vertex2i(x + width, y + height);
            gl::Vertex2i(x + width, y);
            gl::End();
        }
        // Draw a projected latitude-longitude rectangle
        let poly = Polygon::new(&[LatLon{ latitude: Latitude(0.0), longitude: Longitude(0.0) },
                                LatLon{ latitude: Latitude(10.0), longitude: Longitude(0.0) },
                                LatLon{ latitude: Latitude(10.0), longitude: Longitude(10.0) },
                                LatLon{ latitude: Latitude(0.0), longitude: Longitude(10.0) }]);
        let projected = projection.project_poly(&poly);
        unsafe {
            gl::Color3f(0.0, 0.0, 1.0);
            gl::Begin(gl::QUADS);
            for point in projected.points() {
                gl::Vertex2d((x as f64) + point.x, (y as f64) + point.y);
            }
            gl::End();
        }
    }

    fn bounds(&self) -> Option<LatLonRect> {
        None
    }
}
<|endoftext|>


<commit_before>//! The puzzle's state and rules.

use std::collections::BTreeSet;
use std::ops::Index;
use std::rc::Rc;

use ::{Constraint,Val,VarToken};

/// A collection of candidates.
#[derive(Clone,Debug,Eq,PartialEq)]
enum Candidates {
    None,                       // A variable with no candidates.
    Value(Val),                 // A variable set to its initial value.
    Set(Rc<BTreeSet<Val>>),     // A variable with a list of candidates.
}

/// The state of a variable during the solution search.
#[derive(Clone,Debug)]
#[allow(dead_code)]
enum VarState {
    Assigned(Val),
    Unassigned(Candidates),
}

/// The puzzle to be solved.
pub struct Puzzle {
    // The number of variables in the puzzle.
    num_vars: usize,

    // The list of candidates for each variable.
    candidates: Vec<Candidates>,

    // The list of puzzle constraints.
    constraints: Vec<Box<Constraint>>,
}

/// Intermediate puzzle search state.
#[derive(Clone)]
pub struct PuzzleSearch<'a> {
    puzzle: &'a Puzzle,
    vars: Vec<VarState>,
}

/*--------------------------------------------------------------*/

impl Puzzle {
    /// Allocate a new puzzle.
    ///
    /// # Examples
    ///
    /// ```
    /// puzzle_solver::Puzzle::new();
    /// ```
    pub fn new() -> Self {
        Puzzle {
            num_vars: 0,
            candidates: Vec::new(),
            constraints: Vec::new(),
        }
    }

    /// Allocate a new puzzle variable, without inserting any
    /// candidates.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut puzzle = puzzle_solver::Puzzle::new();
    /// puzzle.new_var();
    /// ```
    pub fn new_var(&mut self) -> VarToken {
        let var = VarToken(self.num_vars);
        self.num_vars = self.num_vars + 1;
        self.candidates.push(Candidates::None);
        var
    }

    /// Allocate a new puzzle variable, initialising it with potential
    /// candidates.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut send_more_money = puzzle_solver::Puzzle::new();
    /// send_more_money.new_var_with_candidates(&[0,1,2,3,4,5,6,7,8,9]);
    /// ```
    pub fn new_var_with_candidates(&mut self, candidates: &[Val]) -> VarToken {
        let var = self.new_var();
        self.insert_candidates(var, candidates);
        var
    }

    /// Allocate a 1d vector of puzzle variables, each initialised to
    /// have the same set of potential candidates.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut send_more_money = puzzle_solver::Puzzle::new();
    /// send_more_money.new_vars_with_candidates_1d(8, &[0,1,2,3,4,5,6,7,8,9]);
    /// ```
    pub fn new_vars_with_candidates_1d(&mut self, n: usize, candidates: &[Val])
            -> Vec<VarToken> {
        let mut vars = Vec::with_capacity(n);
        for _ in 0..n {
            vars.push(self.new_var_with_candidates(candidates));
        }
        vars
    }

    /// Allocate a 2d array of puzzle variables, each initialised to
    /// have the same set of potential candidates.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut magic_square = puzzle_solver::Puzzle::new();
    /// magic_square.new_vars_with_candidates_2d(3, 3, &[1,2,3,4,5,6,7,8,9]);
    /// ```
    pub fn new_vars_with_candidates_2d(self: &mut Puzzle,
            width: usize, height: usize, candidates: &[Val])
            -> Vec<Vec<VarToken>> {
        let mut vars = Vec::with_capacity(height);
        for _ in 0..height {
            vars.push(self.new_vars_with_candidates_1d(width, candidates));
        }
        vars
    }

    /// Set a variable to a known value.
    ///
    /// This is useful when the variable is given as part of the
    /// problem.  After this operation, any subsequent attempts to set
    /// the value will panic.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut magic_square = puzzle_solver::Puzzle::new();
    /// let vars = magic_square.new_vars_with_candidates_2d(3, 3,
    ///         &[1,2,3,4,5,6,7,8,9]);
    ///
    /// magic_square.set_value(vars[1][1], 5);
    /// ```
    pub fn set_value(&mut self, var: VarToken, value: Val) {
        let VarToken(idx) = var;

        if let Candidates::Value(val) = self.candidates[idx] {
            if val != value {
                panic!("attempt to set fixed variable");
            }
        }

        self.candidates[idx] = Candidates::Value(value);
    }

    /// Add candidates to a variable.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut send_more_money = puzzle_solver::Puzzle::new();
    /// for _ in 0..9 {
    ///     let var = send_more_money.new_var();
    ///     send_more_money.insert_candidates(var, &[0,1,2,3,4,5,6,7,8,9]);
    /// }
    /// ```
    pub fn insert_candidates(&mut self, var: VarToken, candidates: &[Val]) {
        let VarToken(idx) = var;

        match &self.candidates[idx] {
            &Candidates::Value(_) =>
                panic!("attempt to set fixed variable"),

            &Candidates::None => {
                self.candidates[idx] = Candidates::Set(Rc::new(BTreeSet::new()));
            },

            &Candidates::Set(_) => (),
        }

        // Why you dumb Rust?
        if let Candidates::Set(ref mut rc) = self.candidates[idx] {
            let cs = Rc::get_mut(rc).expect("unique");
            cs.extend(candidates);
        }
    }

    /// Remove candidates from a variable.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut send_more_money = puzzle_solver::Puzzle::new();
    /// let vars = send_more_money.new_vars_with_candidates_1d(8,
    ///         &[0,1,2,3,4,5,6,7,8,9]);
    ///
    /// let s = vars[0];
    /// let m = vars[4];
    /// send_more_money.remove_candidates(s, &[0]);
    /// send_more_money.remove_candidates(m, &[0]);
    /// ```
    pub fn remove_candidates(&mut self, var: VarToken, candidates: &[Val]) {
        let VarToken(idx) = var;

        match self.candidates[idx] {
            Candidates::Value(_) =>
                panic!("attempt to set fixed variable"),

            Candidates::None => (),

            Candidates::Set(ref mut rc) => {
                let cs = Rc::get_mut(rc).expect("unique");
                for c in candidates.iter() {
                    cs.remove(c);
                }
            },
        }
    }

    /// Set the variable's candidates to the intersection with the
    /// given list.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut send_more_money = puzzle_solver::Puzzle::new();
    /// let vars = send_more_money.new_vars_with_candidates_1d(8,
    ///         &[0,1,2,3,4,5,6,7,8,9]);
    ///
    /// let m = vars[4];
    /// send_more_money.intersect_candidates(m, &[0,1]);
    /// ```
    pub fn intersect_candidates(&mut self, var: VarToken, candidates: &[Val]) {
        let VarToken(idx) = var;

        match self.candidates[idx] {
            Candidates::Value(_) =>
                panic!("attempt to set fixed variable"),

            Candidates::None => (),

            Candidates::Set(ref mut rc) => {
                let cs = Rc::get_mut(rc).expect("unique");
                let mut set = BTreeSet::new();
                set.extend(candidates);
                *cs = cs.intersection(&set).cloned().collect();
            },
        }
    }

    /// Add a constraint to the puzzle solution.
    pub fn add_constraint(&mut self, constraint: Box<Constraint>) {
        self.constraints.push(constraint);
    }
}

/*--------------------------------------------------------------*/

impl<'a> PuzzleSearch<'a> {
    /// Allocate a new puzzle searcher.
    #[allow(dead_code)]
    fn new(puzzle: &'a Puzzle) -> Self {
        let mut vars = Vec::with_capacity(puzzle.num_vars);
        for c in puzzle.candidates.iter() {
            vars.push(VarState::Unassigned(c.clone()));
        }

        PuzzleSearch {
            puzzle: puzzle,
            vars: vars,
        }
    }

    /// Check if the variable has been assigned to a value.
    pub fn is_assigned(&self, var: VarToken) -> bool {
        let VarToken(idx) = var;
        match &self.vars[idx] {
            &VarState::Assigned(_) => true,
            &VarState::Unassigned(_) => false,
        }
    }

    /// Get the value assigned to a variable, or None.
    ///
    /// This should be used if the variable may potentially be
    /// unassigned.  For example, when implementing constraints.
    pub fn get_assigned(&self, var: VarToken) -> Option<Val> {
        let VarToken(idx) = var;
        match &self.vars[idx] {
            &VarState::Assigned(val) => Some(val),
            &VarState::Unassigned(_) => None,
        }
    }

    /// Remove a single candidate from an unassigned variable.
    pub fn remove_candidate(&mut self, var: VarToken, val: Val) {
        let VarToken(idx) = var;
        if let VarState::Unassigned(ref mut cs) = self.vars[idx] {
            match cs {
                &mut Candidates::None => return,
                &mut Candidates::Value(v) => {
                    if v == val {
                        *cs = Candidates::None;
                    }
                },
                &mut Candidates::Set(ref mut rc) => {
                    if rc.contains(&val) {
                        let mut set = Rc::make_mut(rc);
                        set.remove(&val);
                    }
                },
            }
        }
    }
}

impl<'a> Index<VarToken> for PuzzleSearch<'a> {
    type Output = Val;

    /// Get the value assigned to a variable.
    ///
    /// # Panics
    ///
    /// Panics if the variable has not been assigned.
    fn index(&self, var: VarToken) -> &Val {
        let VarToken(idx) = var;
        match &self.vars[idx] {
            &VarState::Assigned(ref val) => val,
            &VarState::Unassigned(_) => panic!("unassigned"),
        }
    }
}
<commit_msg>Add puzzle solver stubs.<commit_after>//! The puzzle's state and rules.

use std::collections::BTreeSet;
use std::iter;
use std::iter::Iterator;
use std::ops::Index;
use std::rc::Rc;

use ::{Constraint,Solution,Val,VarToken};

/// A collection of candidates.
#[derive(Clone,Debug,Eq,PartialEq)]
enum Candidates {
    None,                       // A variable with no candidates.
    Value(Val),                 // A variable set to its initial value.
    Set(Rc<BTreeSet<Val>>),     // A variable with a list of candidates.
}

/// The state of a variable during the solution search.
#[derive(Clone,Debug)]
enum VarState {
    Assigned(Val),
    Unassigned(Candidates),
}

/// The puzzle to be solved.
pub struct Puzzle {
    // The number of variables in the puzzle.
    num_vars: usize,

    // The list of candidates for each variable.
    candidates: Vec<Candidates>,

    // The list of puzzle constraints.
    constraints: Vec<Box<Constraint>>,
}

/// Intermediate puzzle search state.
#[derive(Clone)]
pub struct PuzzleSearch<'a> {
    puzzle: &'a Puzzle,
    vars: Vec<VarState>,
}

/*--------------------------------------------------------------*/

impl Candidates {
    /// Count the number of candidates for a variable.
    fn len(&self) -> usize {
        match self {
            &Candidates::None => 0,
            &Candidates::Value(_) => 1,
            &Candidates::Set(ref rc) => rc.len(),
        }
    }

    /// Get an iterator over all of the candidates of a variable.
    fn iter<'a>(&'a self) -> Box<Iterator<Item=Val> + 'a> {
        match self {
            &Candidates::None => Box::new(iter::empty()),
            &Candidates::Value(val) => Box::new(iter::once(val)),
            &Candidates::Set(ref rc) => Box::new(rc.iter().map(|x| *x)),
        }
    }
}

/*--------------------------------------------------------------*/

impl Puzzle {
    /// Allocate a new puzzle.
    ///
    /// # Examples
    ///
    /// ```
    /// puzzle_solver::Puzzle::new();
    /// ```
    pub fn new() -> Self {
        Puzzle {
            num_vars: 0,
            candidates: Vec::new(),
            constraints: Vec::new(),
        }
    }

    /// Allocate a new puzzle variable, without inserting any
    /// candidates.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut puzzle = puzzle_solver::Puzzle::new();
    /// puzzle.new_var();
    /// ```
    pub fn new_var(&mut self) -> VarToken {
        let var = VarToken(self.num_vars);
        self.num_vars = self.num_vars + 1;
        self.candidates.push(Candidates::None);
        var
    }

    /// Allocate a new puzzle variable, initialising it with potential
    /// candidates.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut send_more_money = puzzle_solver::Puzzle::new();
    /// send_more_money.new_var_with_candidates(&[0,1,2,3,4,5,6,7,8,9]);
    /// ```
    pub fn new_var_with_candidates(&mut self, candidates: &[Val]) -> VarToken {
        let var = self.new_var();
        self.insert_candidates(var, candidates);
        var
    }

    /// Allocate a 1d vector of puzzle variables, each initialised to
    /// have the same set of potential candidates.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut send_more_money = puzzle_solver::Puzzle::new();
    /// send_more_money.new_vars_with_candidates_1d(8, &[0,1,2,3,4,5,6,7,8,9]);
    /// ```
    pub fn new_vars_with_candidates_1d(&mut self, n: usize, candidates: &[Val])
            -> Vec<VarToken> {
        let mut vars = Vec::with_capacity(n);
        for _ in 0..n {
            vars.push(self.new_var_with_candidates(candidates));
        }
        vars
    }

    /// Allocate a 2d array of puzzle variables, each initialised to
    /// have the same set of potential candidates.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut magic_square = puzzle_solver::Puzzle::new();
    /// magic_square.new_vars_with_candidates_2d(3, 3, &[1,2,3,4,5,6,7,8,9]);
    /// ```
    pub fn new_vars_with_candidates_2d(self: &mut Puzzle,
            width: usize, height: usize, candidates: &[Val])
            -> Vec<Vec<VarToken>> {
        let mut vars = Vec::with_capacity(height);
        for _ in 0..height {
            vars.push(self.new_vars_with_candidates_1d(width, candidates));
        }
        vars
    }

    /// Set a variable to a known value.
    ///
    /// This is useful when the variable is given as part of the
    /// problem.  After this operation, any subsequent attempts to set
    /// the value will panic.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut magic_square = puzzle_solver::Puzzle::new();
    /// let vars = magic_square.new_vars_with_candidates_2d(3, 3,
    ///         &[1,2,3,4,5,6,7,8,9]);
    ///
    /// magic_square.set_value(vars[1][1], 5);
    /// ```
    pub fn set_value(&mut self, var: VarToken, value: Val) {
        let VarToken(idx) = var;

        if let Candidates::Value(val) = self.candidates[idx] {
            if val != value {
                panic!("attempt to set fixed variable");
            }
        }

        self.candidates[idx] = Candidates::Value(value);
    }

    /// Add candidates to a variable.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut send_more_money = puzzle_solver::Puzzle::new();
    /// for _ in 0..9 {
    ///     let var = send_more_money.new_var();
    ///     send_more_money.insert_candidates(var, &[0,1,2,3,4,5,6,7,8,9]);
    /// }
    /// ```
    pub fn insert_candidates(&mut self, var: VarToken, candidates: &[Val]) {
        let VarToken(idx) = var;

        match &self.candidates[idx] {
            &Candidates::Value(_) =>
                panic!("attempt to set fixed variable"),

            &Candidates::None => {
                self.candidates[idx] = Candidates::Set(Rc::new(BTreeSet::new()));
            },

            &Candidates::Set(_) => (),
        }

        // Why you dumb Rust?
        if let Candidates::Set(ref mut rc) = self.candidates[idx] {
            let cs = Rc::get_mut(rc).expect("unique");
            cs.extend(candidates);
        }
    }

    /// Remove candidates from a variable.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut send_more_money = puzzle_solver::Puzzle::new();
    /// let vars = send_more_money.new_vars_with_candidates_1d(8,
    ///         &[0,1,2,3,4,5,6,7,8,9]);
    ///
    /// let s = vars[0];
    /// let m = vars[4];
    /// send_more_money.remove_candidates(s, &[0]);
    /// send_more_money.remove_candidates(m, &[0]);
    /// ```
    pub fn remove_candidates(&mut self, var: VarToken, candidates: &[Val]) {
        let VarToken(idx) = var;

        match self.candidates[idx] {
            Candidates::Value(_) =>
                panic!("attempt to set fixed variable"),

            Candidates::None => (),

            Candidates::Set(ref mut rc) => {
                let cs = Rc::get_mut(rc).expect("unique");
                for c in candidates.iter() {
                    cs.remove(c);
                }
            },
        }
    }

    /// Set the variable's candidates to the intersection with the
    /// given list.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut send_more_money = puzzle_solver::Puzzle::new();
    /// let vars = send_more_money.new_vars_with_candidates_1d(8,
    ///         &[0,1,2,3,4,5,6,7,8,9]);
    ///
    /// let m = vars[4];
    /// send_more_money.intersect_candidates(m, &[0,1]);
    /// ```
    pub fn intersect_candidates(&mut self, var: VarToken, candidates: &[Val]) {
        let VarToken(idx) = var;

        match self.candidates[idx] {
            Candidates::Value(_) =>
                panic!("attempt to set fixed variable"),

            Candidates::None => (),

            Candidates::Set(ref mut rc) => {
                let cs = Rc::get_mut(rc).expect("unique");
                let mut set = BTreeSet::new();
                set.extend(candidates);
                *cs = cs.intersection(&set).cloned().collect();
            },
        }
    }

    /// Add a constraint to the puzzle solution.
    pub fn add_constraint(&mut self, constraint: Box<Constraint>) {
        self.constraints.push(constraint);
    }

    /// Find any solution to the given puzzle.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut puzzle = puzzle_solver::Puzzle::new();
    /// puzzle.new_var_with_candidates(&[1,2]);
    /// puzzle.new_var_with_candidates(&[3,4]);
    ///
    /// let solution = puzzle.solve_any();
    /// assert!(solution.is_some());
    /// ```
    pub fn solve_any(&self) -> Option<Solution> {
        let mut search = PuzzleSearch::new(self);
        let mut solutions = Vec::with_capacity(1);
        search.solve(1, &mut solutions);
        solutions.pop()
    }

    /// Find the solution to the given puzzle, verifying that it is
    /// unique.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut puzzle = puzzle_solver::Puzzle::new();
    /// puzzle.new_var_with_candidates(&[1,2]);
    /// puzzle.new_var_with_candidates(&[3,4]);
    ///
    /// let solution = puzzle.solve_unique();
    /// assert!(solution.is_none());
    /// ```
    pub fn solve_unique(&self) -> Option<Solution> {
        let mut search = PuzzleSearch::new(self);
        let mut solutions = Vec::with_capacity(2);
        search.solve(2, &mut solutions);
        if solutions.len() == 1 {
            return solutions.pop();
        }

        None
    }

    /// Find all solutions to the given puzzle.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut puzzle = puzzle_solver::Puzzle::new();
    /// puzzle.new_var_with_candidates(&[1,2]);
    /// puzzle.new_var_with_candidates(&[3,4]);
    ///
    /// let solutions = puzzle.solve_all();
    /// assert_eq!(solutions.len(), 4);
    /// ```
    pub fn solve_all(&self) -> Vec<Solution> {
        let mut search = PuzzleSearch::new(self);
        let mut solutions = Vec::new();
        search.solve(::std::usize::MAX, &mut solutions);
        solutions
    }
}

/*--------------------------------------------------------------*/

impl<'a> PuzzleSearch<'a> {
    /// Allocate a new puzzle searcher.
    fn new(puzzle: &'a Puzzle) -> Self {
        let mut vars = Vec::with_capacity(puzzle.num_vars);
        for c in puzzle.candidates.iter() {
            vars.push(VarState::Unassigned(c.clone()));
        }

        PuzzleSearch {
            puzzle: puzzle,
            vars: vars,
        }
    }

    /// Check if the variable has been assigned to a value.
    pub fn is_assigned(&self, var: VarToken) -> bool {
        let VarToken(idx) = var;
        match &self.vars[idx] {
            &VarState::Assigned(_) => true,
            &VarState::Unassigned(_) => false,
        }
    }

    /// Get the value assigned to a variable, or None.
    ///
    /// This should be used if the variable may potentially be
    /// unassigned.  For example, when implementing constraints.
    pub fn get_assigned(&self, var: VarToken) -> Option<Val> {
        let VarToken(idx) = var;
        match &self.vars[idx] {
            &VarState::Assigned(val) => Some(val),
            &VarState::Unassigned(_) => None,
        }
    }

    /// Remove a single candidate from an unassigned variable.
    pub fn remove_candidate(&mut self, var: VarToken, val: Val) {
        let VarToken(idx) = var;
        if let VarState::Unassigned(ref mut cs) = self.vars[idx] {
            match cs {
                &mut Candidates::None => return,
                &mut Candidates::Value(v) => {
                    if v == val {
                        *cs = Candidates::None;
                    }
                },
                &mut Candidates::Set(ref mut rc) => {
                    if rc.contains(&val) {
                        let mut set = Rc::make_mut(rc);
                        set.remove(&val);
                    }
                },
            }
        }
    }

    /// Solve the puzzle, finding up to count solutions.
    fn solve(&mut self, count: usize, solutions: &mut Vec<Solution>) {
        let next_unassigned = self.vars.iter().enumerate().min_by_key(
                |&(_, vs)| match vs {
                    &VarState::Assigned(_) => ::std::usize::MAX,
                    &VarState::Unassigned(ref cs) => cs.len(),
                });

        if let Some((idx, &VarState::Unassigned(ref cs))) = next_unassigned {
            if cs.len() == 0 {
                // Contradiction.
                return;
            }

            for val in cs.iter() {
                let mut new = self.clone();
                new.vars[idx] = VarState::Assigned(val);

                new.solve(count, solutions);
                if solutions.len() >= count {
                    // Reached desired number of solutions.
                    return;
                }
            }
        } else {
            // No unassigned variables remaining.
            let mut vars = Vec::with_capacity(self.puzzle.num_vars);
            for var in self.vars.iter() {
                match var {
                    &VarState::Assigned(val) => vars.push(val),
                    &VarState::Unassigned(_) => unreachable!(),
                }
            }
            solutions.push(Solution{ vars: vars });
        }
    }
}

impl<'a> Index<VarToken> for PuzzleSearch<'a> {
    type Output = Val;

    /// Get the value assigned to a variable.
    ///
    /// # Panics
    ///
    /// Panics if the variable has not been assigned.
    fn index(&self, var: VarToken) -> &Val {
        let VarToken(idx) = var;
        match &self.vars[idx] {
            &VarState::Assigned(ref val) => val,
            &VarState::Unassigned(_) => panic!("unassigned"),
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 Developers of the Rand project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Implementation for SGX using RDRAND instruction
use crate::Error;
use core::arch::x86_64::_rdrand64_step;
use core::mem;
use core::num::NonZeroU32;

// Recommendation from "Intel® Digital Random Number Generator (DRNG) Software
// Implementation Guide" - Section 5.2.1 and "Intel® 64 and IA-32 Architectures
// Software Developer’s Manual" - Volume 1 - Section 7.3.17.1.
const RETRY_LIMIT: usize = 10;
const WORD_SIZE: usize = mem::size_of::<u64>();

#[target_feature(enable = "rdrand")]
unsafe fn rdrand() -> Result<[u8; WORD_SIZE], Error> {
    for _ in 0..RETRY_LIMIT {
        let mut el = mem::uninitialized();
        if _rdrand64_step(&mut el) == 1 {
            // AMD CPUs from families 14h to 16h (pre Ryzen) will sometimes give
            // bogus random data. Discard these values and warn the user.
            // See https://github.com/systemd/systemd/issues/11810#issuecomment-489727505
            if cfg!(not(target_env = "sgx")) && (el == 0 || el == !0) {
                error!("RDRAND returned suspicious value {}, CPU RNG is broken", el);
                return Err(Error::UNKNOWN);
            }
            return Ok(el.to_ne_bytes());
        }
    }
    error!("RDRAND failed, CPU issue likely");
    Err(Error::UNKNOWN)
}

// "rdrand" target feature requires "+rdrnd" flag, see https://github.com/rust-lang/rust/issues/49653.
#[cfg(all(target_env = "sgx", not(target_feature = "rdrand")))]
compile_error!(
    "SGX targets require 'rdrand' target feature. Enable by using -C target-feature=+rdrnd."
);

#[cfg(target_feature = "rdrand")]
fn is_rdrand_supported() -> bool {
    true
}

// TODO use is_x86_feature_detected!("rdrand") when that works in core. See:
// https://github.com/rust-lang-nursery/stdsimd/issues/464
#[cfg(not(target_feature = "rdrand"))]
fn is_rdrand_supported() -> bool {
    use core::arch::x86_64::__cpuid;
    use lazy_static::lazy_static;
    // SAFETY: All x86_64 CPUs support CPUID leaf 1
    const FLAG: u32 = 1 << 30;
    lazy_static! {
        static ref HAS_RDRAND: bool = unsafe { __cpuid(1).ecx & FLAG != 0 };
    }
    *HAS_RDRAND
}

pub fn getrandom_inner(dest: &mut [u8]) -> Result<(), Error> {
    if !is_rdrand_supported() {
        return Err(Error::UNAVAILABLE);
    }

    // SAFETY: After this point, rdrand is supported, so calling the rdrand
    // functions is not undefined behavior.
    unsafe { rdrand_exact(dest) }
}

#[target_feature(enable = "rdrand")]
unsafe fn rdrand_exact(dest: &mut [u8]) -> Result<(), Error> {
    // We use chunks_exact_mut instead of chunks_mut as it allows almost all
    // calls to memcpy to be elided by the compiler.
    let mut chunks = dest.chunks_exact_mut(WORD_SIZE);
    for chunk in chunks.by_ref() {
        chunk.copy_from_slice(&rdrand()?);
    }

    let tail = chunks.into_remainder();
    let n = tail.len();
    if n > 0 {
        tail.copy_from_slice(&rdrand()?[..n]);
    }
    Ok(())
}

#[inline(always)]
pub fn error_msg_inner(_: NonZeroU32) -> Option<&'static str> {
    None
}
<commit_msg>Alternative way to detect AMD bug (#48)<commit_after>// Copyright 2018 Developers of the Rand project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Implementation for SGX using RDRAND instruction
use crate::Error;
use core::arch::x86_64::_rdrand64_step;
use core::mem;
use core::num::NonZeroU32;

// Recommendation from "Intel® Digital Random Number Generator (DRNG) Software
// Implementation Guide" - Section 5.2.1 and "Intel® 64 and IA-32 Architectures
// Software Developer’s Manual" - Volume 1 - Section 7.3.17.1.
const RETRY_LIMIT: usize = 10;
const WORD_SIZE: usize = mem::size_of::<u64>();

#[target_feature(enable = "rdrand")]
unsafe fn rdrand() -> Result<[u8; WORD_SIZE], Error> {
    for _ in 0..RETRY_LIMIT {
        let mut el = mem::uninitialized();
        if _rdrand64_step(&mut el) == 1 {
            // AMD CPUs from families 14h to 16h (pre Ryzen) sometimes fail to
            // set CF on bogus random data, so we check these values explictly.
            // See https://github.com/systemd/systemd/issues/11810#issuecomment-489727505
            // We perform this check regardless of target to guard against
            // any implementation that incorrectly fails to set CF.
            if el != 0 && el != !0 {
                return Ok(el.to_ne_bytes());
            }
            error!("RDRAND returned {:X}, CPU RNG may be broken", el);
            // Keep looping in case this was a false positive.
        }
    }
    error!("RDRAND failed, CPU issue likely");
    Err(Error::UNKNOWN)
}

// "rdrand" target feature requires "+rdrnd" flag, see https://github.com/rust-lang/rust/issues/49653.
#[cfg(all(target_env = "sgx", not(target_feature = "rdrand")))]
compile_error!(
    "SGX targets require 'rdrand' target feature. Enable by using -C target-feature=+rdrnd."
);

#[cfg(target_feature = "rdrand")]
fn is_rdrand_supported() -> bool {
    true
}

// TODO use is_x86_feature_detected!("rdrand") when that works in core. See:
// https://github.com/rust-lang-nursery/stdsimd/issues/464
#[cfg(not(target_feature = "rdrand"))]
fn is_rdrand_supported() -> bool {
    use core::arch::x86_64::__cpuid;
    use lazy_static::lazy_static;
    // SAFETY: All x86_64 CPUs support CPUID leaf 1
    const FLAG: u32 = 1 << 30;
    lazy_static! {
        static ref HAS_RDRAND: bool = unsafe { __cpuid(1).ecx & FLAG != 0 };
    }
    *HAS_RDRAND
}

pub fn getrandom_inner(dest: &mut [u8]) -> Result<(), Error> {
    if !is_rdrand_supported() {
        return Err(Error::UNAVAILABLE);
    }

    // SAFETY: After this point, rdrand is supported, so calling the rdrand
    // functions is not undefined behavior.
    unsafe { rdrand_exact(dest) }
}

#[target_feature(enable = "rdrand")]
unsafe fn rdrand_exact(dest: &mut [u8]) -> Result<(), Error> {
    // We use chunks_exact_mut instead of chunks_mut as it allows almost all
    // calls to memcpy to be elided by the compiler.
    let mut chunks = dest.chunks_exact_mut(WORD_SIZE);
    for chunk in chunks.by_ref() {
        chunk.copy_from_slice(&rdrand()?);
    }

    let tail = chunks.into_remainder();
    let n = tail.len();
    if n > 0 {
        tail.copy_from_slice(&rdrand()?[..n]);
    }
    Ok(())
}

#[inline(always)]
pub fn error_msg_inner(_: NonZeroU32) -> Option<&'static str> {
    None
}
<|endoftext|>


<commit_before>use crc32::Crc32Reader;
use types::ZipFile;
use compression::CompressionMethod;
use spec;
use reader_spec;
use result::{ZipResult, ZipError};
use std::io;
use std::cell::RefCell;
use std::collections::HashMap;
use flate2::FlateReader;
use bzip2::reader::BzDecompressor;

/// Wrapper for reading the contents of a ZIP file.
///
/// ```
/// fn doit() -> zip::result::ZipResult<()>
/// {
///     // For demonstration purposes we read from an empty buffer.
///     // Normally a File object would be used.
///     let buf = [0u8, ..128];
///     let mut reader = std::io::BufReader::new(&buf);
///
///     let zip = try!(zip::ZipReader::new(reader));
///
///     for file in zip.files()
///     {
///         println!("Filename: {}", file.file_name);
///         let mut file_reader = try!(zip.read_file(file));
///         let first_byte = try!(file_reader.read_byte());
///         println!("{}", first_byte);
///     }
///     Ok(())
/// }
///
/// println!("Result: {}", doit());
/// ```
pub struct ZipReader<T>
{
    inner: RefCell<T>,
    files: Vec<ZipFile>,
    names_map: HashMap<String, uint>,
}

fn unsupported_zip_error<T>(detail: &'static str) -> ZipResult<T>
{
    Err(ZipError::UnsupportedZipFile(detail))
}

impl<T: Reader+Seek> ZipReader<T>
{
    /// Opens a ZIP file and parses the content headers.
    pub fn new(mut reader: T) -> ZipResult<ZipReader<T>>
    {
        let footer = try!(spec::CentralDirectoryEnd::find_and_parse(&mut reader));

        if footer.disk_number != footer.disk_with_central_directory { return unsupported_zip_error("Support for multi-disk files is not implemented") }

        let directory_start = footer.central_directory_offset as i64;
        let number_of_files = footer.number_of_files_on_this_disk as uint;

        let mut files = Vec::with_capacity(number_of_files);
        let mut names_map = HashMap::new();

        try!(reader.seek(directory_start, io::SeekSet));
        for _ in range(0, number_of_files)
        {
            let file = try!(reader_spec::central_header_to_zip_file(&mut reader));
            names_map.insert(file.file_name.clone(), files.len());
            files.push(file);
        }

        Ok(ZipReader { inner: RefCell::new(reader), files: files, names_map: names_map })
    }

    /// An iterator over the information of all contained files.
    pub fn files(&self) -> ::std::slice::Items<ZipFile>
    {
        self.files.as_slice().iter()
    }

    /// Search for a file entry by name
    pub fn get(&self, name: &str) -> Option<&ZipFile>
    {
        self.names_map.get(name).map(|index| &self.files[*index])
    }

    /// Gets a reader for a contained zipfile.
    ///
    /// May return `ReaderUnavailable` if there is another reader borrowed.
    pub fn read_file(&self, file: &ZipFile) -> ZipResult<Box<Reader>>
    {
        let mut inner_reader = match self.inner.try_borrow_mut()
        {
            Some(reader) => reader,
            None => return Err(ZipError::ReaderUnavailable),
        };
        let pos = file.data_start as i64;

        if file.encrypted
        {
            return unsupported_zip_error("Encrypted files are not supported")
        }

        try!(inner_reader.seek(pos, io::SeekSet));
        let refmut_reader = ::util::RefMutReader::new(inner_reader);
        let limit_reader = io::util::LimitReader::new(refmut_reader, file.compressed_size as uint);

        let reader = match file.compression_method
        {
            CompressionMethod::Stored =>
            {
                box
                    Crc32Reader::new(
                        limit_reader,
                        file.crc32)
                    as Box<Reader>
            },
            CompressionMethod::Deflated =>
            {
                let deflate_reader = limit_reader.deflate_decode();
                box
                    Crc32Reader::new(
                        deflate_reader,
                        file.crc32)
                    as Box<Reader>
            },
            CompressionMethod::Bzip2 =>
            {
                let bzip2_reader = BzDecompressor::new(limit_reader);
                box
                    Crc32Reader::new(
                        bzip2_reader,
                        file.crc32)
                    as Box<Reader>
            },
            _ => return unsupported_zip_error("Compression method not supported"),
        };
        Ok(reader)
    }

    /// Unwrap and return the inner reader object
    ///
    /// The position of the reader is undefined.
    pub fn into_inner(self) -> T
    {
        self.inner.into_inner()
    }

    /// Deprecated method equal to `into_inner()`
    #[deprecated="renamed to into_inner()"]
    pub fn unwrap(self) -> T
    {
        self.into_inner()
    }
}
<commit_msg>std::slice::Items -> Iter<commit_after>use crc32::Crc32Reader;
use types::ZipFile;
use compression::CompressionMethod;
use spec;
use reader_spec;
use result::{ZipResult, ZipError};
use std::io;
use std::cell::RefCell;
use std::collections::HashMap;
use flate2::FlateReader;
use bzip2::reader::BzDecompressor;

/// Wrapper for reading the contents of a ZIP file.
///
/// ```
/// fn doit() -> zip::result::ZipResult<()>
/// {
///     // For demonstration purposes we read from an empty buffer.
///     // Normally a File object would be used.
///     let buf = [0u8, ..128];
///     let mut reader = std::io::BufReader::new(&buf);
///
///     let zip = try!(zip::ZipReader::new(reader));
///
///     for file in zip.files()
///     {
///         println!("Filename: {}", file.file_name);
///         let mut file_reader = try!(zip.read_file(file));
///         let first_byte = try!(file_reader.read_byte());
///         println!("{}", first_byte);
///     }
///     Ok(())
/// }
///
/// println!("Result: {}", doit());
/// ```
pub struct ZipReader<T>
{
    inner: RefCell<T>,
    files: Vec<ZipFile>,
    names_map: HashMap<String, uint>,
}

fn unsupported_zip_error<T>(detail: &'static str) -> ZipResult<T>
{
    Err(ZipError::UnsupportedZipFile(detail))
}

impl<T: Reader+Seek> ZipReader<T>
{
    /// Opens a ZIP file and parses the content headers.
    pub fn new(mut reader: T) -> ZipResult<ZipReader<T>>
    {
        let footer = try!(spec::CentralDirectoryEnd::find_and_parse(&mut reader));

        if footer.disk_number != footer.disk_with_central_directory { return unsupported_zip_error("Support for multi-disk files is not implemented") }

        let directory_start = footer.central_directory_offset as i64;
        let number_of_files = footer.number_of_files_on_this_disk as uint;

        let mut files = Vec::with_capacity(number_of_files);
        let mut names_map = HashMap::new();

        try!(reader.seek(directory_start, io::SeekSet));
        for _ in range(0, number_of_files)
        {
            let file = try!(reader_spec::central_header_to_zip_file(&mut reader));
            names_map.insert(file.file_name.clone(), files.len());
            files.push(file);
        }

        Ok(ZipReader { inner: RefCell::new(reader), files: files, names_map: names_map })
    }

    /// An iterator over the information of all contained files.
    pub fn files(&self) -> ::std::slice::Iter<ZipFile>
    {
        self.files.as_slice().iter()
    }

    /// Search for a file entry by name
    pub fn get(&self, name: &str) -> Option<&ZipFile>
    {
        self.names_map.get(name).map(|index| &self.files[*index])
    }

    /// Gets a reader for a contained zipfile.
    ///
    /// May return `ReaderUnavailable` if there is another reader borrowed.
    pub fn read_file(&self, file: &ZipFile) -> ZipResult<Box<Reader>>
    {
        let mut inner_reader = match self.inner.try_borrow_mut()
        {
            Some(reader) => reader,
            None => return Err(ZipError::ReaderUnavailable),
        };
        let pos = file.data_start as i64;

        if file.encrypted
        {
            return unsupported_zip_error("Encrypted files are not supported")
        }

        try!(inner_reader.seek(pos, io::SeekSet));
        let refmut_reader = ::util::RefMutReader::new(inner_reader);
        let limit_reader = io::util::LimitReader::new(refmut_reader, file.compressed_size as uint);

        let reader = match file.compression_method
        {
            CompressionMethod::Stored =>
            {
                box
                    Crc32Reader::new(
                        limit_reader,
                        file.crc32)
                    as Box<Reader>
            },
            CompressionMethod::Deflated =>
            {
                let deflate_reader = limit_reader.deflate_decode();
                box
                    Crc32Reader::new(
                        deflate_reader,
                        file.crc32)
                    as Box<Reader>
            },
            CompressionMethod::Bzip2 =>
            {
                let bzip2_reader = BzDecompressor::new(limit_reader);
                box
                    Crc32Reader::new(
                        bzip2_reader,
                        file.crc32)
                    as Box<Reader>
            },
            _ => return unsupported_zip_error("Compression method not supported"),
        };
        Ok(reader)
    }

    /// Unwrap and return the inner reader object
    ///
    /// The position of the reader is undefined.
    pub fn into_inner(self) -> T
    {
        self.inner.into_inner()
    }

    /// Deprecated method equal to `into_inner()`
    #[deprecated="renamed to into_inner()"]
    pub fn unwrap(self) -> T
    {
        self.into_inner()
    }
}
<|endoftext|>


<commit_before>//! Logic handling reading from Avro format at user level.
use std::collections::VecDeque;
use std::io::{ErrorKind, Read};
use std::rc::Rc;
use std::str::{from_utf8, FromStr};

use failure::{err_msg, Error};
use serde_json::from_slice;

use decode::decode;
use schema::ParseSchemaError;
use schema::Schema;
use types::Value;
use util::DecodeError;
use Codec;

/// Main interface for reading Avro formatted values.
///
/// To be used as an iterator:
///
/// ```no_run
/// # use avro_rs::Reader;
/// # use std::io::Cursor;
/// # let input = Cursor::new(Vec::<u8>::new());
/// for value in Reader::new(input).unwrap() {
///     match value {
///         Ok(v) => println!("{:?}", v),
///         Err(e) => println!("Error: {}", e),
///     };
/// }
/// ```
pub struct Reader<'a, R> {
    reader: R,
    reader_schema: Option<&'a Schema>,
    writer_schema: Schema,
    codec: Codec,
    marker: [u8; 16],
    items: VecDeque<Value>,
    errored: bool,
}

impl<'a, R: Read> Reader<'a, R> {
    /// Creates a `Reader` given something implementing the `io::Read` trait to read from.
    /// No reader `Schema` will be set.
    ///
    /// **NOTE** The avro header is going to be read automatically upon creation of the `Reader`.
    pub fn new(reader: R) -> Result<Reader<'a, R>, Error> {
        let mut reader = Reader {
            reader,
            reader_schema: None,
            writer_schema: Schema::Null,
            codec: Codec::Null,
            marker: [0u8; 16],
            items: VecDeque::new(),
            errored: false,
        };
        reader.read_header()?;
        Ok(reader)
    }

    /// Creates a `Reader` given a reader `Schema` and something implementing the `io::Read` trait
    /// to read from.
    ///
    /// **NOTE** The avro header is going to be read automatically upon creation of the `Reader`.
    pub fn with_schema(schema: &'a Schema, reader: R) -> Result<Reader<'a, R>, Error> {
        let mut reader = Reader {
            reader,
            reader_schema: Some(schema),
            writer_schema: Schema::Null,
            codec: Codec::Null,
            marker: [0u8; 16],
            items: VecDeque::new(),
            errored: false,
        };
        reader.read_header()?;
        Ok(reader)
    }

    /// Get a reference to the writer `Schema`.
    pub fn writer_schema(&self) -> &Schema {
        &self.writer_schema
    }

    /// Get a reference to the optional reader `Schema`.
    pub fn reader_schema(&self) -> Option<&Schema> {
        self.reader_schema
    }

    /// Try to read the header and to set the writer `Schema`, the `Codec` and the marker based on
    /// its content.
    fn read_header(&mut self) -> Result<(), Error> {
        let meta_schema = Schema::Map(Rc::new(Schema::Bytes));

        let mut buf = [0u8; 4];
        self.reader.read_exact(&mut buf)?;

        if buf != ['O' as u8, 'b' as u8, 'j' as u8, 1u8] {
            return Err(DecodeError::new("wrong magic in header").into())
        }

        if let Value::Map(meta) = decode(&meta_schema, &mut self.reader)? {
            // TODO: surface original parse schema errors instead of coalescing them here
            let schema = meta.get("avro.schema")
                .and_then(|bytes| {
                    if let &Value::Bytes(ref bytes) = bytes {
                        from_slice(bytes.as_ref()).ok()
                    } else {
                        None
                    }
                })
                .and_then(|json| Schema::parse(&json).ok());
            if let Some(schema) = schema {
                self.writer_schema = schema;
            } else {
                return Err(ParseSchemaError::new("unable to parse schema").into())
            }

            if let Some(codec) = meta.get("avro.codec")
                .and_then(|codec| {
                    if let &Value::Bytes(ref bytes) = codec {
                        from_utf8(bytes.as_ref()).ok()
                    } else {
                        None
                    }
                })
                .and_then(|codec| Codec::from_str(codec).ok())
            {
                self.codec = codec;
            }
        } else {
            return Err(DecodeError::new("no metadata in header").into())
        }

        let mut buf = [0u8; 16];
        self.reader.read_exact(&mut buf)?;
        self.marker = buf;

        Ok(())
    }

    /// Try to read a data block, also performing schema resolution for the objects contained in
    /// the block. The objects are stored in an internal buffer to the `Reader`.
    fn read_block(&mut self) -> Result<(), Error> {
        match decode(&Schema::Long, &mut self.reader) {
            Ok(block) => {
                if let Value::Long(block_len) = block {
                    if let Value::Long(block_bytes) = decode(&Schema::Long, &mut self.reader)? {
                        let mut bytes = vec![0u8; block_bytes as usize];
                        self.reader.read_exact(&mut bytes)?;

                        let mut marker = [0u8; 16];
                        self.reader.read_exact(&mut marker)?;

                        if marker != self.marker {
                            return Err(DecodeError::new(
                                "block marker does not match header marker",
                            ).into())
                        }

                        self.codec.decompress(&mut bytes)?;

                        self.items.clear();
                        self.items.reserve_exact(block_len as usize);

                        for _ in 0..block_len {
                            let item = from_avro_datum(
                                &self.writer_schema,
                                &mut &bytes[..],
                                self.reader_schema,
                            )?;
                            self.items.push_back(item)
                        }

                        return Ok(())
                    }
                }
            },
            Err(e) => match e.downcast::<::std::io::Error>()?.kind() {
                // to not return any error in case we only finished to read cleanly from the stream
                ErrorKind::UnexpectedEof => return Ok(()),
                _ => (),
            },
        };
        Err(DecodeError::new("unable to read block").into())
    }
}

impl<'a, R: Read> Iterator for Reader<'a, R> {
    type Item = Result<Value, Error>;

    fn next(&mut self) -> Option<Self::Item> {
        // to prevent keep on reading after the first error occurs
        if self.errored {
            return None
        };
        if self.items.len() == 0 {
            if let Err(e) = self.read_block() {
                self.errored = true;
                return Some(Err(err_msg(e)))
            }
        }

        self.items.pop_front().map(Ok)
    }
}

/// Decode a `Value` encoded in Avro format given its `Schema` and anything implementing `io::Read`
/// to read from.
///
/// In case a reader `Schema` is provided, schema resolution will also be performed.
///
/// **NOTE** This function has a quite small niche of usage and does NOT take care of reading the
/// header and consecutive data blocks; use [`Reader`](struct.Reader.html) if you don't know what
/// you are doing, instead.
pub fn from_avro_datum<R: Read>(
    writer_schema: &Schema,
    reader: &mut R,
    reader_schema: Option<&Schema>,
) -> Result<Value, Error> {
    let value = decode(writer_schema, reader)?;
    match reader_schema {
        Some(ref schema) => value.resolve(schema),
        None => Ok(value),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Cursor;
    use types::{Record, ToAvro};
    use Reader;

    static SCHEMA: &'static str = r#"
            {
                "type": "record",
                "name": "test",
                "fields": [
                    {"name": "a", "type": "long", "default": 42},
                    {"name": "b", "type": "string"}
                ]
            }
        "#;
    static ENCODED: &'static [u8] = &[
        79u8, 98u8, 106u8, 1u8, 4u8, 22u8, 97u8, 118u8, 114u8, 111u8, 46u8, 115u8, 99u8, 104u8,
        101u8, 109u8, 97u8, 222u8, 1u8, 123u8, 34u8, 116u8, 121u8, 112u8, 101u8, 34u8, 58u8, 34u8,
        114u8, 101u8, 99u8, 111u8, 114u8, 100u8, 34u8, 44u8, 34u8, 110u8, 97u8, 109u8, 101u8, 34u8,
        58u8, 34u8, 116u8, 101u8, 115u8, 116u8, 34u8, 44u8, 34u8, 102u8, 105u8, 101u8, 108u8,
        100u8, 115u8, 34u8, 58u8, 91u8, 123u8, 34u8, 110u8, 97u8, 109u8, 101u8, 34u8, 58u8, 34u8,
        97u8, 34u8, 44u8, 34u8, 116u8, 121u8, 112u8, 101u8, 34u8, 58u8, 34u8, 108u8, 111u8, 110u8,
        103u8, 34u8, 44u8, 34u8, 100u8, 101u8, 102u8, 97u8, 117u8, 108u8, 116u8, 34u8, 58u8, 52u8,
        50u8, 125u8, 44u8, 123u8, 34u8, 110u8, 97u8, 109u8, 101u8, 34u8, 58u8, 34u8, 98u8, 34u8,
        44u8, 34u8, 116u8, 121u8, 112u8, 101u8, 34u8, 58u8, 34u8, 115u8, 116u8, 114u8, 105u8,
        110u8, 103u8, 34u8, 125u8, 93u8, 125u8, 20u8, 97u8, 118u8, 114u8, 111u8, 46u8, 99u8, 111u8,
        100u8, 101u8, 99u8, 8u8, 110u8, 117u8, 108u8, 108u8, 0u8, 94u8, 61u8, 54u8, 221u8, 190u8,
        207u8, 108u8, 180u8, 158u8, 57u8, 114u8, 40u8, 173u8, 199u8, 228u8, 239u8, 4u8, 20u8, 54u8,
        6u8, 102u8, 111u8, 111u8, 54u8, 6u8, 102u8, 111u8, 111u8, 94u8, 61u8, 54u8, 221u8, 190u8,
        207u8, 108u8, 180u8, 158u8, 57u8, 114u8, 40u8, 173u8, 199u8, 228u8, 239,
    ];

    #[test]
    fn test_from_avro_datum() {
        let schema = Schema::parse_str(SCHEMA).unwrap();
        let mut encoded: &'static [u8] = &[54, 6, 102, 111, 111];

        let mut record = Record::new(&schema).unwrap();
        record.put("a", 27i64);
        record.put("b", "foo");
        let expected = record.avro();

        assert_eq!(
            from_avro_datum(&schema, &mut encoded, None).unwrap(),
            expected
        );
    }

    #[test]
    fn test_reader_iterator() {
        let schema = Schema::parse_str(SCHEMA).unwrap();
        let reader = Reader::with_schema(&schema, ENCODED).unwrap();

        let mut record = Record::new(&schema).unwrap();
        record.put("a", 27i64);
        record.put("b", "foo");
        let expected = record.avro();

        for value in reader {
            assert_eq!(value.unwrap(), expected);
        }
    }

    #[test]
    fn test_reader_invalid_header() {
        let schema = Schema::parse_str(SCHEMA).unwrap();
        let invalid = ENCODED.to_owned().into_iter().skip(1).collect::<Vec<u8>>();
        assert!(Reader::with_schema(&schema, &invalid[..]).is_err());
    }

    #[test]
    fn test_reader_invalid_block() {
        let schema = Schema::parse_str(SCHEMA).unwrap();
        let invalid = ENCODED
            .to_owned()
            .into_iter()
            .rev()
            .skip(19)
            .collect::<Vec<u8>>()
            .into_iter()
            .rev()
            .collect::<Vec<u8>>();
        let reader = Reader::with_schema(&schema, &invalid[..]).unwrap();
        for value in reader {
            assert!(value.is_err());
        }
    }

    #[test]
    fn test_reader_empty_buffer() {
        let empty = Cursor::new(Vec::new());
        Reader::new(empty).is_err();
    }

    #[test]
    fn test_reader_only_header() {
        let invalid = ENCODED
            .to_owned()
            .into_iter()
            .take(165)
            .collect::<Vec<u8>>();
        let reader = Reader::new(&invalid[..]).unwrap();
        for value in reader {
            assert!(value.is_err());
        }
    }
}
<commit_msg>Use a single &[u8] when decoding items in Reader (#32)<commit_after>//! Logic handling reading from Avro format at user level.
use std::collections::VecDeque;
use std::io::{ErrorKind, Read};
use std::rc::Rc;
use std::str::{from_utf8, FromStr};

use failure::{err_msg, Error};
use serde_json::from_slice;

use decode::decode;
use schema::ParseSchemaError;
use schema::Schema;
use types::Value;
use util::DecodeError;
use Codec;

/// Main interface for reading Avro formatted values.
///
/// To be used as an iterator:
///
/// ```no_run
/// # use avro_rs::Reader;
/// # use std::io::Cursor;
/// # let input = Cursor::new(Vec::<u8>::new());
/// for value in Reader::new(input).unwrap() {
///     match value {
///         Ok(v) => println!("{:?}", v),
///         Err(e) => println!("Error: {}", e),
///     };
/// }
/// ```
pub struct Reader<'a, R> {
    reader: R,
    reader_schema: Option<&'a Schema>,
    writer_schema: Schema,
    codec: Codec,
    marker: [u8; 16],
    items: VecDeque<Value>,
    errored: bool,
}

impl<'a, R: Read> Reader<'a, R> {
    /// Creates a `Reader` given something implementing the `io::Read` trait to read from.
    /// No reader `Schema` will be set.
    ///
    /// **NOTE** The avro header is going to be read automatically upon creation of the `Reader`.
    pub fn new(reader: R) -> Result<Reader<'a, R>, Error> {
        let mut reader = Reader {
            reader,
            reader_schema: None,
            writer_schema: Schema::Null,
            codec: Codec::Null,
            marker: [0u8; 16],
            items: VecDeque::new(),
            errored: false,
        };
        reader.read_header()?;
        Ok(reader)
    }

    /// Creates a `Reader` given a reader `Schema` and something implementing the `io::Read` trait
    /// to read from.
    ///
    /// **NOTE** The avro header is going to be read automatically upon creation of the `Reader`.
    pub fn with_schema(schema: &'a Schema, reader: R) -> Result<Reader<'a, R>, Error> {
        let mut reader = Reader {
            reader,
            reader_schema: Some(schema),
            writer_schema: Schema::Null,
            codec: Codec::Null,
            marker: [0u8; 16],
            items: VecDeque::new(),
            errored: false,
        };
        reader.read_header()?;
        Ok(reader)
    }

    /// Get a reference to the writer `Schema`.
    pub fn writer_schema(&self) -> &Schema {
        &self.writer_schema
    }

    /// Get a reference to the optional reader `Schema`.
    pub fn reader_schema(&self) -> Option<&Schema> {
        self.reader_schema
    }

    /// Try to read the header and to set the writer `Schema`, the `Codec` and the marker based on
    /// its content.
    fn read_header(&mut self) -> Result<(), Error> {
        let meta_schema = Schema::Map(Rc::new(Schema::Bytes));

        let mut buf = [0u8; 4];
        self.reader.read_exact(&mut buf)?;

        if buf != ['O' as u8, 'b' as u8, 'j' as u8, 1u8] {
            return Err(DecodeError::new("wrong magic in header").into())
        }

        if let Value::Map(meta) = decode(&meta_schema, &mut self.reader)? {
            // TODO: surface original parse schema errors instead of coalescing them here
            let schema = meta.get("avro.schema")
                .and_then(|bytes| {
                    if let &Value::Bytes(ref bytes) = bytes {
                        from_slice(bytes.as_ref()).ok()
                    } else {
                        None
                    }
                })
                .and_then(|json| Schema::parse(&json).ok());
            if let Some(schema) = schema {
                self.writer_schema = schema;
            } else {
                return Err(ParseSchemaError::new("unable to parse schema").into())
            }

            if let Some(codec) = meta.get("avro.codec")
                .and_then(|codec| {
                    if let &Value::Bytes(ref bytes) = codec {
                        from_utf8(bytes.as_ref()).ok()
                    } else {
                        None
                    }
                })
                .and_then(|codec| Codec::from_str(codec).ok())
            {
                self.codec = codec;
            }
        } else {
            return Err(DecodeError::new("no metadata in header").into())
        }

        let mut buf = [0u8; 16];
        self.reader.read_exact(&mut buf)?;
        self.marker = buf;

        Ok(())
    }

    /// Try to read a data block, also performing schema resolution for the objects contained in
    /// the block. The objects are stored in an internal buffer to the `Reader`.
    fn read_block(&mut self) -> Result<(), Error> {
        match decode(&Schema::Long, &mut self.reader) {
            Ok(block) => {
                if let Value::Long(block_len) = block {
                    if let Value::Long(block_bytes) = decode(&Schema::Long, &mut self.reader)? {
                        let mut bytes = vec![0u8; block_bytes as usize];
                        self.reader.read_exact(&mut bytes)?;

                        let mut marker = [0u8; 16];
                        self.reader.read_exact(&mut marker)?;

                        if marker != self.marker {
                            return Err(DecodeError::new(
                                "block marker does not match header marker",
                            ).into())
                        }

                        self.codec.decompress(&mut bytes)?;

                        self.items.clear();
                        self.items.reserve_exact(block_len as usize);

                        let mut block_bytes = &bytes[..];
                        for _ in 0..block_len {
                            let item = from_avro_datum(
                                &self.writer_schema,
                                &mut block_bytes,
                                self.reader_schema,
                            )?;
                            self.items.push_back(item)
                        }

                        return Ok(())
                    }
                }
            },
            Err(e) => match e.downcast::<::std::io::Error>()?.kind() {
                // to not return any error in case we only finished to read cleanly from the stream
                ErrorKind::UnexpectedEof => return Ok(()),
                _ => (),
            },
        };
        Err(DecodeError::new("unable to read block").into())
    }
}

impl<'a, R: Read> Iterator for Reader<'a, R> {
    type Item = Result<Value, Error>;

    fn next(&mut self) -> Option<Self::Item> {
        // to prevent keep on reading after the first error occurs
        if self.errored {
            return None
        };
        if self.items.len() == 0 {
            if let Err(e) = self.read_block() {
                self.errored = true;
                return Some(Err(err_msg(e)))
            }
        }

        self.items.pop_front().map(Ok)
    }
}

/// Decode a `Value` encoded in Avro format given its `Schema` and anything implementing `io::Read`
/// to read from.
///
/// In case a reader `Schema` is provided, schema resolution will also be performed.
///
/// **NOTE** This function has a quite small niche of usage and does NOT take care of reading the
/// header and consecutive data blocks; use [`Reader`](struct.Reader.html) if you don't know what
/// you are doing, instead.
pub fn from_avro_datum<R: Read>(
    writer_schema: &Schema,
    reader: &mut R,
    reader_schema: Option<&Schema>,
) -> Result<Value, Error> {
    let value = decode(writer_schema, reader)?;
    match reader_schema {
        Some(ref schema) => value.resolve(schema),
        None => Ok(value),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Cursor;
    use types::{Record, ToAvro};
    use Reader;

    static SCHEMA: &'static str = r#"
            {
                "type": "record",
                "name": "test",
                "fields": [
                    {"name": "a", "type": "long", "default": 42},
                    {"name": "b", "type": "string"}
                ]
            }
        "#;
    static ENCODED: &'static [u8] = &[
        79u8, 98u8, 106u8, 1u8, 4u8, 22u8, 97u8, 118u8, 114u8, 111u8, 46u8, 115u8, 99u8, 104u8,
        101u8, 109u8, 97u8, 222u8, 1u8, 123u8, 34u8, 116u8, 121u8, 112u8, 101u8, 34u8, 58u8, 34u8,
        114u8, 101u8, 99u8, 111u8, 114u8, 100u8, 34u8, 44u8, 34u8, 110u8, 97u8, 109u8, 101u8, 34u8,
        58u8, 34u8, 116u8, 101u8, 115u8, 116u8, 34u8, 44u8, 34u8, 102u8, 105u8, 101u8, 108u8,
        100u8, 115u8, 34u8, 58u8, 91u8, 123u8, 34u8, 110u8, 97u8, 109u8, 101u8, 34u8, 58u8, 34u8,
        97u8, 34u8, 44u8, 34u8, 116u8, 121u8, 112u8, 101u8, 34u8, 58u8, 34u8, 108u8, 111u8, 110u8,
        103u8, 34u8, 44u8, 34u8, 100u8, 101u8, 102u8, 97u8, 117u8, 108u8, 116u8, 34u8, 58u8, 52u8,
        50u8, 125u8, 44u8, 123u8, 34u8, 110u8, 97u8, 109u8, 101u8, 34u8, 58u8, 34u8, 98u8, 34u8,
        44u8, 34u8, 116u8, 121u8, 112u8, 101u8, 34u8, 58u8, 34u8, 115u8, 116u8, 114u8, 105u8,
        110u8, 103u8, 34u8, 125u8, 93u8, 125u8, 20u8, 97u8, 118u8, 114u8, 111u8, 46u8, 99u8, 111u8,
        100u8, 101u8, 99u8, 8u8, 110u8, 117u8, 108u8, 108u8, 0u8, 94u8, 61u8, 54u8, 221u8, 190u8,
        207u8, 108u8, 180u8, 158u8, 57u8, 114u8, 40u8, 173u8, 199u8, 228u8, 239u8, 4u8, 20u8, 54u8,
        6u8, 102u8, 111u8, 111u8, 84u8, 6u8, 98u8, 97u8, 114u8, 94u8, 61u8, 54u8, 221u8, 190u8,
        207u8, 108u8, 180u8, 158u8, 57u8, 114u8, 40u8, 173u8, 199u8, 228u8, 239u8,
    ];

    #[test]
    fn test_from_avro_datum() {
        let schema = Schema::parse_str(SCHEMA).unwrap();
        let mut encoded: &'static [u8] = &[54, 6, 102, 111, 111];

        let mut record = Record::new(&schema).unwrap();
        record.put("a", 27i64);
        record.put("b", "foo");
        let expected = record.avro();

        assert_eq!(
            from_avro_datum(&schema, &mut encoded, None).unwrap(),
            expected
        );
    }

    #[test]
    fn test_reader_iterator() {
        let schema = Schema::parse_str(SCHEMA).unwrap();
        let reader = Reader::with_schema(&schema, ENCODED).unwrap();

        let mut record1 = Record::new(&schema).unwrap();
        record1.put("a", 27i64);
        record1.put("b", "foo");

        let mut record2 = Record::new(&schema).unwrap();
        record2.put("a", 42i64);
        record2.put("b", "bar");

        let expected = vec![record1.avro(), record2.avro()];

        for (i, value) in reader.enumerate() {
            assert_eq!(value.unwrap(), expected[i]);
        }
    }

    #[test]
    fn test_reader_invalid_header() {
        let schema = Schema::parse_str(SCHEMA).unwrap();
        let invalid = ENCODED.to_owned().into_iter().skip(1).collect::<Vec<u8>>();
        assert!(Reader::with_schema(&schema, &invalid[..]).is_err());
    }

    #[test]
    fn test_reader_invalid_block() {
        let schema = Schema::parse_str(SCHEMA).unwrap();
        let invalid = ENCODED
            .to_owned()
            .into_iter()
            .rev()
            .skip(19)
            .collect::<Vec<u8>>()
            .into_iter()
            .rev()
            .collect::<Vec<u8>>();
        let reader = Reader::with_schema(&schema, &invalid[..]).unwrap();
        for value in reader {
            assert!(value.is_err());
        }
    }

    #[test]
    fn test_reader_empty_buffer() {
        let empty = Cursor::new(Vec::new());
        Reader::new(empty).is_err();
    }

    #[test]
    fn test_reader_only_header() {
        let invalid = ENCODED
            .to_owned()
            .into_iter()
            .take(165)
            .collect::<Vec<u8>>();
        let reader = Reader::new(&invalid[..]).unwrap();
        for value in reader {
            assert!(value.is_err());
        }
    }
}
<|endoftext|>