<commit_before>//! The parser is based on the following grammar
//!
//! Equation ==> Term { "+" Term }
//!           |  Term { "-" Term }
//!
//! Term     ==> Function Product
//!           |  Product
//!
//! Product  ==> Factor { "*" Term }
//!           |  Factor { "/" Term }
//!
//! Factor   ==> "-" Factor
//!           |  Exponent { "^" Factor }
//!
//! Exponent ==> Number { "!" }
//!
//! Number   ==> "(" Equation ")"
//!           |  "|" Equation "|"
//!           |  Constant
//!           |  Digit { Digit } | Digit { Digit } "." { Digit }
//!
//! Function ==> "sin" | "cos" | "tan" | "asin" | "acos" | "atan" | "sqrt" | "abs" | "exp"
//!           |  "ln" | "log"
//!
//! Constant ==> "pi" | "e" | "phi"
//!
//! Digit    ==> "0".."9"

use std::vec::IntoIter;
use std::iter::Peekable;
use errors::{CalcrResult, CalcrError};
use ast::Ast;
use ast::AstVal;
use ast::AstBranch::*;
use ast::FuncKind::*;
use ast::ConstKind::*;
use ast::OpKind::*;
use token::Token;
use token::TokVal;
use token::TokVal::*;

pub fn parse_equation(tokens: Vec<Token>) -> CalcrResult<Ast> {
    let end_pos = tokens.last().and_then(|tok| Some(tok.span.1)).unwrap_or(0);
    let mut parser = Parser {
        iter: tokens.into_iter().peekable(),
        paren_level: 0,
        abs_level: 0,
        end_pos: end_pos,
    };
    parser.parse_equation()
}

pub struct Parser {
    iter: Peekable<IntoIter<Token>>,
    paren_level: u32,
    abs_level: u32,
    end_pos: usize,
}

impl Parser {
    fn parse_equation(&mut self) -> CalcrResult<Ast> {
        let mut lhs = try!(self.parse_term());
        while self.peek_tok_val().map_or(false, |val| *val == Op(Plus) || *val == Op(Minus)) {
            let Token { val: tok_val, span: tok_span } = self.consume_tok();
            let rhs = try!(self.parse_term());
            lhs = Ast {
                val: AstVal::Op(tok_val.op().unwrap()),
                span: tok_span,
                branches: Binary(Box::new(lhs), Box::new(rhs)),
            }
        }
        if self.peek_tok_val().map_or(false, |val| *val == ParenClose) && self.paren_level < 1 {
            let Token { val: _, span: tok_span } = self.consume_tok();
            Err(CalcrError {
                desc: format!("Missing opening parentheses"),
                span: Some(tok_span),
            })
        } else if self.peek_tok_val().map_or(false, |val| *val == AbsDelim) && self.abs_level < 1 {
            let Token { val: _, span: tok_span } = self.consume_tok();
            Err(CalcrError {
                desc: format!("Missing opening abs delimiter"),
                span: Some(tok_span),
            })
        } else {
            Ok(lhs)
        }
    }

    fn parse_term(&mut self) -> CalcrResult<Ast> {
        // check if we have a function
        let func_opt = self.peek_tok_val().and_then(|val| {
            if let Name(ref name) = *val {
                match name.as_ref() {
                    "cos" => Some(Cos),
                    "sin" => Some(Sin),
                    "tan" => Some(Tan),
                    "asin" => Some(Asin),
                    "acos" => Some(Acos),
                    "atan" => Some(Atan),
                    "sqrt" => Some(Sqrt),
                    "abs" => Some(Abs),
                    "exp" => Some(Exp),
                    "ln" => Some(Ln),
                    "log" => Some(Log),
                    _ => None,
                }
            } else {
                None
            }
        });

        if let Some(func) = func_opt {
            let Token { val: _, span: tok_span } = self.consume_tok();
            let arg = try!(self.parse_product());
            Ok(Ast {
                val: AstVal::Func(func),
                span: tok_span,
                branches: Unary(Box::new(arg)),
            })
        } else {
            self.parse_product()
        }
    }

    fn parse_product(&mut self) -> CalcrResult<Ast> {
        let mut lhs = try!(self.parse_factor());
        while self.peek_tok_val().map_or(false, |val| *val == Op(Mult) || *val == Op(Div)) {
            let Token { val: tok_val, span: tok_span } = self.consume_tok();
            let rhs = try!(self.parse_term());
            lhs = Ast {
                val: AstVal::Op(tok_val.op().unwrap()),
                span: tok_span,
                branches: Binary(Box::new(lhs), Box::new(rhs)),
            };
        }
        Ok(lhs)
    }

    fn parse_factor(&mut self) -> CalcrResult<Ast> {
        // when we lex we only store `Minus`s since we do not have any context there,
        // however we know if we see a `Minus` now, then it is a `Neg`.
        if self.peek_tok_val().map_or(false, |val| *val == Op(Minus)) {
            let tok_span = self.consume_tok().span;
            let rhs = try!(self.parse_factor());
            Ok(Ast {
                val: AstVal::Op(Neg),
                span: tok_span,
                branches: Unary(Box::new(rhs)),
            })
        } else {
            let lhs = try!(self.parse_exponent());
            if self.peek_tok_val().map_or(false, |val| *val == Op(Pow)) {
                let tok_span = self.consume_tok().span;
                let rhs = try!(self.parse_factor());
                Ok(Ast {
                    val: AstVal::Op(Pow),
                    span: tok_span,
                    branches: Binary(Box::new(lhs), Box::new(rhs)),
                })
            } else {
                Ok(lhs)
            }
        }
    }

    fn parse_exponent(&mut self) -> CalcrResult<Ast> {
        let mut out = try!(self.parse_number());

        while self.peek_tok_val().map_or(false, |val| *val == Op(Fact)) {
            let tok_span = self.consume_tok().span;
            out = Ast {
                val: AstVal::Op(Fact),
                span: tok_span,
                branches: Unary(Box::new(out)),
            };
        }
        Ok(out)
    }

    fn parse_number(&mut self) -> CalcrResult<Ast> {
        if self.toks_empty() {
            Err(CalcrError {
                desc: format!("Expected number or constant"),
                span: Some((self.end_pos, self.end_pos)),
            })
        } else {
            let Token { val: tok_val, span: tok_span } = self.consume_tok();
            match tok_val {
                ParenOpen => {
                    self.paren_level += 1;
                    let eq = try!(self.parse_equation());
                    if self.peek_tok_val().map_or(false, |val| *val != ParenClose) {
                        Err(CalcrError {
                            desc: "Missing closing parentheses".to_string(),
                            span: Some(tok_span),
                        })
                    } else {
                        self.paren_level -= 1;
                        let close_paren_span = self.consume_tok().span;
                        Ok(Ast {
                            val: AstVal::Paren,
                            span: (tok_span.0, close_paren_span.1),
                            branches: Unary(Box::new(eq)),
                        })
                    }
                },
                AbsDelim => {
                    self.abs_level += 1;
                    let eq = try!(self.parse_equation());
                    if self.peek_tok_val().map_or(false, |val| *val != AbsDelim) {
                        Err(CalcrError {
                            desc: "Missing closing abs delimiter".to_string(),
                            span: Some(tok_span),
                        })
                    } else {
                        self.abs_level -= 1;
                        let close_delim_span = self.consume_tok().span;
                        Ok(Ast {
                            val: AstVal::Func(Abs),
                            span: (tok_span.0, close_delim_span.1),
                            branches: Unary(Box::new(eq)),
                        })
                    }
                },
                Name(ref name) => {
                    // at this point any Name, HAS to be a known constant
                    let cnst = match name.as_ref() {
                        "pi" => Pi,
                        "e" => E,
                        "phi" => Phi,
                        _ => return Err(CalcrError {
                            desc: format!("Invalid function or constant: {}", name),
                            span: Some(tok_span),
                        }),
                    };
                    Ok(Ast {
                        val: AstVal::Const(cnst),
                        span: tok_span,
                        branches: Leaf,
                    })
                },
                Num(num) => {
                    Ok(Ast {
                        val: AstVal::Num(num),
                        span: tok_span,
                        branches: Leaf,
                    })
                },
                _ => Err(CalcrError {
                    desc: format!("Expected number or constant"),
                    span: Some(tok_span),
                }),
            }
        }
    }

    /// Checks if we have run out of `Token`s to parse
    fn toks_empty(&mut self) -> bool {
        self.iter.peek().is_none()
    }

    /// Peeks at the next `Token` and returns `Some` if one was found, or `None` if none are left
    fn peek_tok_val(&mut self) -> Option<&TokVal> {
        self.iter.peek().and_then(|ref tok| Some(&tok.val))
    }

    /// Consumes a `Token` - thereby advanding `pos` - and returns it
    ///
    /// # Panics
    /// This function panics if there are no more `Token`s to consume
    fn consume_tok(&mut self) -> Token {
        let tok = self.iter.next();
        tok.unwrap()
    }
}<commit_msg>Change how the parser parses function<commit_after>//! The parser is based on the following grammar
//!
//! Equation ==> Product { "+" Product }
//!           |  Product { "-" Product }
//!
//! Product  ==> Factor { "*" Factor }
//!           |  Factor { "/" Factor }
//!
//! Factor   ==> "-" Factor
//!           |  Exponent { "^" Factor }
//!
//! Exponent ==> Number { "!" }
//!
//! Number   ==> Function "(" Equation ")"
//!           |  Constant
//!           |  "(" Equation ")"
//!           |  "|" Equation "|"
//!           |  Digit { Digit } | Digit { Digit } "." { Digit }
//!
//! Function ==> "sin" | "cos" | "tan" | "asin" | "acos" | "atan" | "sqrt" | "abs" | "exp"
//!           |  "ln" | "log"
//!
//! Constant ==> "pi" | "e" | "phi"
//!
//! Digit    ==> "0".."9"

use std::vec::IntoIter;
use std::iter::Peekable;
use errors::{CalcrResult, CalcrError};
use ast::Ast;
use ast::AstVal;
use ast::AstBranch::*;
use ast::FuncKind::*;
use ast::ConstKind::*;
use ast::OpKind::*;
use token::Token;
use token::TokVal;
use token::TokVal::*;

pub fn parse_equation(tokens: Vec<Token>) -> CalcrResult<Ast> {
    let end_pos = tokens.last().and_then(|tok| Some(tok.span.1)).unwrap_or(0);
    let mut parser = Parser {
        iter: tokens.into_iter().peekable(),
        paren_level: 0,
        abs_level: 0,
        end_pos: end_pos,
    };
    parser.parse_equation()
}

pub struct Parser {
    iter: Peekable<IntoIter<Token>>,
    paren_level: u32,
    abs_level: u32,
    end_pos: usize,
}

impl Parser {
    fn parse_equation(&mut self) -> CalcrResult<Ast> {
        let mut lhs = try!(self.parse_product());
        while self.peek_tok_val().map_or(false, |val| *val == Op(Plus) || *val == Op(Minus)) {
            let Token { val: tok_val, span: tok_span } = self.consume_tok();
            let rhs = try!(self.parse_product());
            lhs = Ast {
                val: AstVal::Op(tok_val.op().unwrap()),
                span: tok_span,
                branches: Binary(Box::new(lhs), Box::new(rhs)),
            }
        }
        if self.peek_tok_val().map_or(false, |val| *val == ParenClose) && self.paren_level < 1 {
            let Token { val: _, span: tok_span } = self.consume_tok();
            Err(CalcrError {
                desc: format!("Missing opening parentheses"),
                span: Some(tok_span),
            })
        } else if self.peek_tok_val().map_or(false, |val| *val == AbsDelim) && self.abs_level < 1 {
            let Token { val: _, span: tok_span } = self.consume_tok();
            Err(CalcrError {
                desc: format!("Missing opening abs delimiter"),
                span: Some(tok_span),
            })
        } else {
            Ok(lhs)
        }
    }

    fn parse_product(&mut self) -> CalcrResult<Ast> {
        let mut lhs = try!(self.parse_factor());
        while self.peek_tok_val().map_or(false, |val| *val == Op(Mult) || *val == Op(Div)) {
            let Token { val: tok_val, span: tok_span } = self.consume_tok();
            let rhs = try!(self.parse_factor());
            lhs = Ast {
                val: AstVal::Op(tok_val.op().unwrap()),
                span: tok_span,
                branches: Binary(Box::new(lhs), Box::new(rhs)),
            };
        }
        Ok(lhs)
    }

    fn parse_factor(&mut self) -> CalcrResult<Ast> {
        // when we lex we only store `Minus`s since we do not have any context there,
        // however we know if we see a `Minus` now, then it is a `Neg`.
        if self.peek_tok_val().map_or(false, |val| *val == Op(Minus)) {
            let tok_span = self.consume_tok().span;
            let rhs = try!(self.parse_factor());
            Ok(Ast {
                val: AstVal::Op(Neg),
                span: tok_span,
                branches: Unary(Box::new(rhs)),
            })
        } else {
            let lhs = try!(self.parse_exponent());
            if self.peek_tok_val().map_or(false, |val| *val == Op(Pow)) {
                let tok_span = self.consume_tok().span;
                let rhs = try!(self.parse_factor());
                Ok(Ast {
                    val: AstVal::Op(Pow),
                    span: tok_span,
                    branches: Binary(Box::new(lhs), Box::new(rhs)),
                })
            } else {
                Ok(lhs)
            }
        }
    }

    fn parse_exponent(&mut self) -> CalcrResult<Ast> {
        let mut out = try!(self.parse_number());

        while self.peek_tok_val().map_or(false, |val| *val == Op(Fact)) {
            let tok_span = self.consume_tok().span;
            out = Ast {
                val: AstVal::Op(Fact),
                span: tok_span,
                branches: Unary(Box::new(out)),
            };
        }
        Ok(out)
    }

    fn parse_number(&mut self) -> CalcrResult<Ast> {
        if self.toks_empty() {
            Err(CalcrError {
                desc: format!("Expected number or constant"),
                span: Some((self.end_pos, self.end_pos)),
            })
        } else {
            let Token { val: tok_val, span: tok_span } = self.consume_tok();
            match tok_val {
                Name(ref name) => {
                    let val = match name.as_ref() {
                        "pi" => AstVal::Const(Pi),
                        "e" => AstVal::Const(E),
                        "phi" => AstVal::Const(Phi),
                        "cos" => AstVal::Func(Cos),
                        "sin" => AstVal::Func(Sin),
                        "tan" => AstVal::Func(Tan),
                        "asin" => AstVal::Func(Asin),
                        "acos" => AstVal::Func(Acos),
                        "atan" => AstVal::Func(Atan),
                        "sqrt" => AstVal::Func(Sqrt),
                        "abs" => AstVal::Func(Abs),
                        "exp" => AstVal::Func(Exp),
                        "ln" => AstVal::Func(Ln),
                        "log" => AstVal::Func(Log),
                        _ => return Err(CalcrError {
                            desc: format!("Invalid function or constant: {}", name),
                            span: Some(tok_span),
                        }),
                    };
                    if let AstVal::Func(_) = val {
                        // it's a function so we need to grab its argument
                        if self.peek_tok_val().map_or(false, |val| *val == ParenOpen) {
                            // since we know the next token is an open paren, we use
                            // this function to get its AST
                            let arg = try!(self.parse_equation());
                            Ok(Ast {
                                val: val,
                                span: tok_span,
                                branches: Unary(Box::new(arg)),
                            })
                        } else {
                            Err(CalcrError {
                                desc: "Missing opening parentheses after function".to_string(),
                                span: Some(tok_span),
                            })
                        }
                    } else {
                        Ok(Ast {
                            val: val,
                            span: tok_span,
                            branches: Leaf,
                        })
                    }
                },
                ParenOpen => {
                    self.paren_level += 1;
                    let eq = try!(self.parse_equation());
                    if self.peek_tok_val().map_or(false, |val| *val != ParenClose) {
                        Err(CalcrError {
                            desc: "Missing closing parentheses".to_string(),
                            span: Some(tok_span),
                        })
                    } else {
                        self.paren_level -= 1;
                        let close_paren_span = self.consume_tok().span;
                        Ok(Ast {
                            val: AstVal::Paren,
                            span: (tok_span.0, close_paren_span.1),
                            branches: Unary(Box::new(eq)),
                        })
                    }
                },
                AbsDelim => {
                    self.abs_level += 1;
                    let eq = try!(self.parse_equation());
                    if self.peek_tok_val().map_or(false, |val| *val != AbsDelim) {
                        Err(CalcrError {
                            desc: "Missing closing abs delimiter".to_string(),
                            span: Some(tok_span),
                        })
                    } else {
                        self.abs_level -= 1;
                        let close_delim_span = self.consume_tok().span;
                        Ok(Ast {
                            val: AstVal::Func(Abs),
                            span: (tok_span.0, close_delim_span.1),
                            branches: Unary(Box::new(eq)),
                        })
                    }
                },
                Num(num) => {
                    Ok(Ast {
                        val: AstVal::Num(num),
                        span: tok_span,
                        branches: Leaf,
                    })
                },
                _ => Err(CalcrError {
                    desc: format!("Expected number or constant"),
                    span: Some(tok_span),
                }),
            }
        }
    }

    /// Checks if we have run out of `Token`s to parse
    fn toks_empty(&mut self) -> bool {
        self.iter.peek().is_none()
    }

    /// Peeks at the next `Token` and returns `Some` if one was found, or `None` if none are left
    fn peek_tok_val(&mut self) -> Option<&TokVal> {
        self.iter.peek().and_then(|ref tok| Some(&tok.val))
    }

    /// Consumes a `Token` - thereby advanding `pos` - and returns it
    ///
    /// # Panics
    /// This function panics if there are no more `Token`s to consume
    fn consume_tok(&mut self) -> Token {
        let tok = self.iter.next();
        tok.unwrap()
    }
}<|endoftext|>


<commit_before>use super::models::{LanguageModel, WordVec};
use self::Token::*;
use super::mayberef::MaybeRef::{self, Val, Ref};
use std::fmt::{Display, Formatter, Error};

pub fn parse(expr: &str, model: &LanguageModel) -> Result<WordVec, String> {
    expression(&mut Tokens::from(expr.chars()), model).map(|w| w.take())
}

#[derive(PartialEq, Eq, Debug, Clone)]
enum Token {
    RParen,
    LParen,
    Plus,
    Minus,
    Word(String),
    Invalid(char),
}

impl Display for Token {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        let invalid;
        fmt.write_str(match *self {
            RParen => ")",
            LParen => "(",
            Plus => "+",
            Minus => "-",
            Word(ref w) => &*w,
            Invalid(w) => {
                invalid = w.to_string();
                &*invalid
            },
        })
    }
}

fn expression<'a, I>(tokens: &mut Tokens<I>, model: &'a LanguageModel)
    -> Result<MaybeRef<'a, WordVec>, String>
    where I: Iterator<Item = char>
{
    let token = try!(tokens.next().ok_or("An expression may not be empty"));
    match token {
        LParen => {
            let vec = try!(expression(tokens, model));
            let rparen = try!(tokens.next().ok_or("Unbalanced parentheses"));
            if rparen != RParen {
                return Err(format!("Expected ')', found '{}'", rparen))
            }
            rhs(vec, tokens, model)
        },
        Word(word) => {
            let vec = try!(model.get(&*word).ok_or_else(||
                format!("'{}' is not present in the language model", &word)
            ));
            println!("{:?}", vec);
            rhs(Ref(vec), tokens, model)
        },
        _ => Err(format!("An expression may not start with '{}'", token)),
    }
}

fn rhs<'a, I>(lhs: MaybeRef<'a, WordVec>, tokens: &mut Tokens<I>, model: &LanguageModel)
    -> Result<MaybeRef<'a, WordVec>, String>
    where I: Iterator<Item = char>
{
    match tokens.peek() {
        Some(&RParen) | None => return Ok(lhs),
        _ => {}
    }

    let token = tokens.next().unwrap();
    match token {
        Plus => Ok(Val(lhs.take() + &*try!(expression(tokens, model)))),
        Minus => Ok(Val(lhs.take() - &*try!(expression(tokens, model)))),
        Word(_) | LParen => Err(format!("'{}' found in invalid position", token)),
        Invalid(s) => Err(format!("'{}' could not be tokenized", s)),
        RParen => unreachable!(),
    }
}

struct Tokens<I> {
    iter: I,
    next_char: Option<Option<char>>,
    next_token: Option<Option<Token>>
}

impl<I> Tokens<I> where I: Iterator<Item = char> {
    fn from(iter: I) -> Tokens<I> {
        Tokens {
            iter: iter,
            next_char: None,
            next_token: None,
        }
    }

    fn get_next_token(&mut self) -> Option<Token> {
        while let Some(' ') = self.peek_char() {
            self.take();
        }

        match self.peek_char() {
            Some('a'...'z') => return Some(self.word()),
            // Some('0'..'9') | Some('.') => return Some(self.number()),
            None => return None,
            _ => {}
        }

        Some(match self.take().unwrap() {
            '+' => Plus,
            '-' => Minus,
            '(' => LParen,
            ')' => RParen,
            c => Invalid(c),
        })
    }

    fn peek(&mut self) -> Option<&Token> {
        if let Some(ref t) = self.next_token {
            return t.as_ref();
        }

        self.next_token = Some(self.get_next_token());
        return self.next_token.as_ref().unwrap().as_ref();
    }

    fn peek_char(&mut self) -> Option<char> {
        match self.next_char {
            Some(x) => x,
            None => {
                let next = self.iter.next();
                self.next_char = Some(next);
                next
            }
        }
    }

    fn take(&mut self) -> Option<char> {
        self.peek_char();
        self.next_char.take().unwrap()
    }

    fn word(&mut self) -> Token {
        let mut token = String::new();
        while let Some('a'...'z') = self.peek_char() {
            token.push(self.take().unwrap());
        }

        Word(token)
    }

}

impl<I> Iterator for Tokens<I> where I: Iterator<Item = char> {
    type Item = Token;

    fn next(&mut self) -> Option<Token> {
        self.peek();
        self.next_token.take().unwrap()
    }
}

mod test {
    use super::Tokens;
    use super::Token::*;
    use super::expression;
    use super::super::models::{WordVec, LanguageModel, LanguageModelBuilder};

    #[test]
    fn single_character() {
        let c = Tokens::from("a".chars()).next().unwrap();
        assert_eq!(c, Word("a".to_string()));
    }

    #[test]
    fn single_token() {
        let s = Tokens::from("abc".chars()).next().unwrap();
        assert_eq!(s, Word("abc".to_string()));
    }

    #[test]
    fn two_tokens() {
        let s = Tokens::from("abc def".chars()).collect::<Vec<_>>();
        assert_eq!(s, vec![Word("abc".to_string()), Word("def".to_string())]);
    }

    #[test]
    fn plus_tokens() {
        let s = Tokens::from("   abc  +def".chars()).collect::<Vec<_>>();
        assert_eq!(s, vec![Word("abc".to_string()), Plus, Word("def".to_string())]);
    }

    #[test]
    fn bracket_tokens() {
        let s = Tokens::from(" ()  (abc  +)def".chars()).collect::<Vec<_>>();
        assert_eq!(s, vec![LParen, RParen, LParen, Word("abc".to_string()),
            Plus, RParen, Word("def".to_string())]);
    }

    #[test]
    fn test_associativity() {
        let words = vec!("a".to_string(), "b".to_string(), "c".to_string());
        let mut builder = LanguageModelBuilder::new(1, words);

        {
            let mut acc = builder.new_file();
            let text = "x x a a b a c b x x";
            for word in text.split(' ') {
                acc.add_word(word.to_string());
            }
        }

        let model = builder.build();
        let mut a = Tokens::from("a - b + c".chars());
        let mut b = Tokens::from("(a - b) + c".chars());
        assert_eq!(expression(&mut a, &model).unwrap().take(),
                   expression(&mut b, &model).unwrap().take());
    }
}<commit_msg>Fix warning in test<commit_after>use super::models::{LanguageModel, WordVec};
use self::Token::*;
use super::mayberef::MaybeRef::{self, Val, Ref};
use std::fmt::{Display, Formatter, Error};

pub fn parse(expr: &str, model: &LanguageModel) -> Result<WordVec, String> {
    expression(&mut Tokens::from(expr.chars()), model).map(|w| w.take())
}

#[derive(PartialEq, Eq, Debug, Clone)]
enum Token {
    RParen,
    LParen,
    Plus,
    Minus,
    Word(String),
    Invalid(char),
}

impl Display for Token {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        let invalid;
        fmt.write_str(match *self {
            RParen => ")",
            LParen => "(",
            Plus => "+",
            Minus => "-",
            Word(ref w) => &*w,
            Invalid(w) => {
                invalid = w.to_string();
                &*invalid
            },
        })
    }
}

fn expression<'a, I>(tokens: &mut Tokens<I>, model: &'a LanguageModel)
    -> Result<MaybeRef<'a, WordVec>, String>
    where I: Iterator<Item = char>
{
    let token = try!(tokens.next().ok_or("An expression may not be empty"));
    match token {
        LParen => {
            let vec = try!(expression(tokens, model));
            let rparen = try!(tokens.next().ok_or("Unbalanced parentheses"));
            if rparen != RParen {
                return Err(format!("Expected ')', found '{}'", rparen))
            }
            rhs(vec, tokens, model)
        },
        Word(word) => {
            let vec = try!(model.get(&*word).ok_or_else(||
                format!("'{}' is not present in the language model", &word)
            ));
            println!("{:?}", vec);
            rhs(Ref(vec), tokens, model)
        },
        _ => Err(format!("An expression may not start with '{}'", token)),
    }
}

fn rhs<'a, I>(lhs: MaybeRef<'a, WordVec>, tokens: &mut Tokens<I>, model: &LanguageModel)
    -> Result<MaybeRef<'a, WordVec>, String>
    where I: Iterator<Item = char>
{
    match tokens.peek() {
        Some(&RParen) | None => return Ok(lhs),
        _ => {}
    }

    let token = tokens.next().unwrap();
    match token {
        Plus => Ok(Val(lhs.take() + &*try!(expression(tokens, model)))),
        Minus => Ok(Val(lhs.take() - &*try!(expression(tokens, model)))),
        Word(_) | LParen => Err(format!("'{}' found in invalid position", token)),
        Invalid(s) => Err(format!("'{}' could not be tokenized", s)),
        RParen => unreachable!(),
    }
}

struct Tokens<I> {
    iter: I,
    next_char: Option<Option<char>>,
    next_token: Option<Option<Token>>
}

impl<I> Tokens<I> where I: Iterator<Item = char> {
    fn from(iter: I) -> Tokens<I> {
        Tokens {
            iter: iter,
            next_char: None,
            next_token: None,
        }
    }

    fn get_next_token(&mut self) -> Option<Token> {
        while let Some(' ') = self.peek_char() {
            self.take();
        }

        match self.peek_char() {
            Some('a'...'z') => return Some(self.word()),
            // Some('0'..'9') | Some('.') => return Some(self.number()),
            None => return None,
            _ => {}
        }

        Some(match self.take().unwrap() {
            '+' => Plus,
            '-' => Minus,
            '(' => LParen,
            ')' => RParen,
            c => Invalid(c),
        })
    }

    fn peek(&mut self) -> Option<&Token> {
        if let Some(ref t) = self.next_token {
            return t.as_ref();
        }

        self.next_token = Some(self.get_next_token());
        return self.next_token.as_ref().unwrap().as_ref();
    }

    fn peek_char(&mut self) -> Option<char> {
        match self.next_char {
            Some(x) => x,
            None => {
                let next = self.iter.next();
                self.next_char = Some(next);
                next
            }
        }
    }

    fn take(&mut self) -> Option<char> {
        self.peek_char();
        self.next_char.take().unwrap()
    }

    fn word(&mut self) -> Token {
        let mut token = String::new();
        while let Some('a'...'z') = self.peek_char() {
            token.push(self.take().unwrap());
        }

        Word(token)
    }

}

impl<I> Iterator for Tokens<I> where I: Iterator<Item = char> {
    type Item = Token;

    fn next(&mut self) -> Option<Token> {
        self.peek();
        self.next_token.take().unwrap()
    }
}

mod test {

    #[test]
    fn single_character() {
        use super::Tokens;
        let c = Tokens::from("a".chars()).next().unwrap();
        assert_eq!(c, Word("a".to_string()));
    }

    #[test]
    fn single_token() {
        use super::Tokens;
        let s = Tokens::from("abc".chars()).next().unwrap();
        assert_eq!(s, Word("abc".to_string()));
    }

    #[test]
    fn two_tokens() {
        use super::Tokens;
        let s = Tokens::from("abc def".chars()).collect::<Vec<_>>();
        assert_eq!(s, vec![Word("abc".to_string()), Word("def".to_string())]);
    }

    #[test]
    fn plus_tokens() {
        use super::Tokens;
        let s = Tokens::from("   abc  +def".chars()).collect::<Vec<_>>();
        assert_eq!(s, vec![Word("abc".to_string()), Plus, Word("def".to_string())]);
    }

    #[test]
    fn bracket_tokens() {
        use super::Tokens;
        use super::Token::*;
        let s = Tokens::from(" ()  (abc  +)def".chars()).collect::<Vec<_>>();
        assert_eq!(s, vec![LParen, RParen, LParen, Word("abc".to_string()),
            Plus, RParen, Word("def".to_string())]);
    }

    #[test]
    fn test_associativity() {
        use super::expression;
        use super::super::models::{WordVec, LanguageModel, LanguageModelBuilder};

        let words = vec!("a".to_string(), "b".to_string(), "c".to_string());
        let mut builder = LanguageModelBuilder::new(1, words);

        {
            let mut acc = builder.new_file();
            let text = "x x a a b a c b x x";
            for word in text.split(' ') {
                acc.add_word(word.to_string());
            }
        }

        let model = builder.build();
        let mut a = Tokens::from("a - b + c".chars());
        let mut b = Tokens::from("(a - b) + c".chars());
        assert_eq!(expression(&mut a, &model).unwrap().take(),
                   expression(&mut b, &model).unwrap().take());
    }
}<|endoftext|>


<commit_before>use codemap::Spanned;
use lexer::Token;
use std::iter::Peekable;
use std::slice::Iter;

type ParseIter<'a> = Peekable<Iter<'a, Spanned<Token>>>;

#[derive(Debug, Clone, PartialEq)]
pub enum Node {
  Block(Vec<Node>),
  If{cond: Box<Node>},
  Stmt(Box<Node>),
  Break,
  Continue,
  Expr,
  Pass,

  // Literals
  Null,
  Bool(bool),
  Float(f64),
  Int(i64),
  Str(String),
  Name(String),

  None,
}

#[derive(Debug, Clone, PartialEq)]
pub enum ParseErrorKind {
  UnexpectedToken(Token),
  UnexpectedEOF,
}

#[derive(Debug, Clone, PartialEq)]
pub enum Parse {
  Ok(Node),
  Err(ParseErrorKind),
  None
}

impl Parse {
  fn expect(self, msg: &str) -> Node {
    match self {
      Parse::Ok(n) => n,
      Parse::Err(x) => panic!("{:?} from {:?}", msg, x),
      Parse::None => panic!("{:?}", msg),
    }
  }

  fn map<F>(self, op: F) -> Parse
      where F: FnOnce(Node) -> Node {
    if let Parse::Ok(n) = self {
      Parse::Ok(op(n))
    }
    else {
      self
    }
  }
}

// Return true if the next token in `it` is `kind`
fn peek_token(it: &mut ParseIter, kind: Token) -> bool {
  if let Some(&c) = it.peek() {
    c.node == kind
  }
  else {
    false
  }
}

// Return true if the next token in `it` is `kind` *and* consume the token
fn use_token(it: &mut ParseIter, kind: Token) -> bool {
  if let Some(&c) = it.peek() {
    if c.node == kind {
      it.next();
    }
    c.node == kind
  }
  else {
    false
  }
}

// Panic if the next token in `it` is *not* `kind`
fn require_token(it: &mut ParseIter, kind: Token) {
  if let Some(&c) = it.peek() {
    if c.node != kind {
      panic!("Expected {:?} but found {:?}", kind, c.node);
    }
    it.next();
  } else {
    panic!("Expected {:?} but found end of stream", kind);
  }
}

fn parse_expr(it: &mut ParseIter) -> Parse {
  if let Some(&c) = it.peek() {
    return match c.node {
      Token::Null => {it.next(); Parse::Ok(Node::Null)},
      Token::Bool(x) => {it.next(); Parse::Ok(Node::Bool(x))},
      Token::Float(x) => {it.next(); Parse::Ok(Node::Float(x))},
      Token::Int(x) => {it.next(); Parse::Ok(Node::Int(x))},
      Token::Str(ref x) => {it.next(); Parse::Ok(Node::Str(x.clone()))},
      Token::Name(ref x) => {it.next(); Parse::Ok(Node::Name(x.clone()))},
      _ => Parse::None,
    }
  }
  Parse::None
}

fn parse_stmt(it: &mut ParseIter) -> Parse {
  if let Some(&c) = it.peek() {
    match c.node {
      Token::Break => {
        it.next();
        Parse::Ok(Node::Break)
      }

      Token::Continue => {
        it.next();
        Parse::Ok(Node::Continue)
      }

      Token::Pass => {
        it.next();
        Parse::Ok(Node::Pass)
      }

      Token::If => {
        it.next();
        parse_expr(it).map(|expr| Node::If{cond: Box::new(expr)})
      }

      _ => {
        parse_expr(it).map(|expr| Node::Stmt(Box::new(expr)))
      }
    }
  }
  else {
    Parse::None
  }
}

pub fn parse(tokens: Vec<Spanned<Token>>) -> Parse {
  let mut it: ParseIter = tokens.iter().peekable();
  let mut nodes: Vec<Node> = vec![];
  loop {
    match parse_stmt(&mut it) {
      Parse::Ok(node) => nodes.push(node),
      Parse::Err(err) => {return Parse::Err(err)},
      Parse::None => {break},
    }

    require_token(&mut it, Token::Newline);
    if peek_token(&mut it, Token::EOF) {
      break;
    }
  }
  Parse::Ok(Node::Block(nodes))
}
<commit_msg>parser: Switch to Result instead of my bullshit<commit_after>use codemap::Spanned;
use lexer::Token;
use std::iter::Peekable;
use std::slice::Iter;
use self::ParseErrorKind::*;

type ParseIter<'a> = Peekable<Iter<'a, Spanned<Token>>>;
type Parse = Result<Node, ParseErrorKind>;

#[derive(Debug, Clone, PartialEq)]
pub enum Node {
  Block(Vec<Node>),
  If{cond: Box<Node>},
  Stmt(Box<Node>),
  Break,
  Continue,
  Expr,
  Pass,

  // Literals
  Null,
  Bool(bool),
  Float(f64),
  Int(i64),
  Str(String),
  Name(String),
}

#[derive(Debug, Clone, PartialEq)]
pub enum ParseErrorKind {
  UnexpectedToken(Token),
  UnexpectedEOF,
}

// Return true if the next token in `it` is `kind`
fn peek_token(it: &mut ParseIter, kind: Token) -> bool {
  if let Some(&c) = it.peek() {
    c.node == kind
  }
  else {
    false
  }
}

// Return true if the next token in `it` is `kind` *and* consume the token
fn use_token(it: &mut ParseIter, kind: Token) -> bool {
  if let Some(&c) = it.peek() {
    if c.node == kind {
      it.next();
    }
    c.node == kind
  }
  else {
    false
  }
}

// Panic if the next token in `it` is *not* `kind`
fn require_token(it: &mut ParseIter, kind: Token) -> Result<(), ParseErrorKind> {
  if let Some(&c) = it.peek() {
    it.next();

    if c.node == kind {
      return Ok(());
    }

    return Err(UnexpectedToken(c.node.clone()));
  }

  return Err(UnexpectedEOF);
}

fn parse_expr(it: &mut ParseIter) -> Parse {
  if let Some(&c) = it.peek() {
    return match c.node {
      Token::Null => {it.next(); Ok(Node::Null)},
      Token::Bool(x) => {it.next(); Ok(Node::Bool(x))},
      Token::Float(x) => {it.next(); Ok(Node::Float(x))},
      Token::Int(x) => {it.next(); Ok(Node::Int(x))},
      Token::Str(ref x) => {it.next(); Ok(Node::Str(x.clone()))},
      Token::Name(ref x) => {it.next(); Ok(Node::Name(x.clone()))},
      _ => Err(UnexpectedToken(c.node.clone())),
    }
  }

  Err(UnexpectedEOF)
}

fn parse_stmt(it: &mut ParseIter) -> Parse {
  if let Some(&c) = it.peek() {
    return match c.node {
      Token::Break => {
        it.next();
        Ok(Node::Break)
      }

      Token::Continue => {
        it.next();
        Ok(Node::Continue)
      }

      Token::Pass => {
        it.next();
        Ok(Node::Pass)
      }

      Token::If => {
        it.next();
        parse_expr(it).map(|expr| Node::If{cond: Box::new(expr)})
      }

      _ => {
        parse_expr(it).map(|expr| Node::Stmt(Box::new(expr)))
      }
    }
  }

  Err(UnexpectedEOF)
}

pub fn parse(tokens: Vec<Spanned<Token>>) -> Parse {
  let mut it: ParseIter = tokens.iter().peekable();
  let mut nodes: Vec<Node> = vec![];

  while !peek_token(&mut it, Token::EOF) {
    nodes.push(parse_stmt(&mut it)?);
    require_token(&mut it, Token::Newline)?;
  }

  Ok(Node::Block(nodes))
}
<|endoftext|>


<commit_before>//! Postgres-based Fact Database
//!
//! # Design Notes
//!
//! ## Scope
//!
//! The general philsophy is that things having to do with persistence go here,
//! while things related to non-persistent components go in `holmes::engine`.
//!
//! In the long run, we would like to persist nearly everything in the
//! database, so that a client-server model can one bay restored. However,
//! in the short term this has little benefit, so only the items needing to
//! use SQL for efficient execution are being included.
//!
//! The biggest hurdle here is the persistence of code:
//!
//! * How do we store Types?
//! * How do we store bound functions?
//!
//! One possible long term answer is Cap'n' Proto `SturdyRef`s
//!
//! ## Other Databases
//!
//! For the moment, this is the only implementation, and there are no others
//! on the horizon, so this interface is not abstract.
//!
//! The only major hurdle to using another backend would be figuring out how
//! to make the `dyn` module abstract over databases.
use std::collections::hash_map::HashMap;
use std::collections::hash_map::Entry::{Occupied, Vacant};

use postgres::{rows, Connection, SslMode};
use postgres::types::{FromSql, ToSql};

use engine::types::{Fact, Predicate, MatchExpr, Clause};

mod error;
pub mod dyn;

pub use self::error::Error;
use self::dyn::types;
use self::dyn::{Type, Value};
use fact_db::{FactDB, Result};

/// An iterator over a `postgres::rows::Row`.
/// It does not implement the normal iter interface because it does not have
/// a set item type, but it implements a similar interface for ease of use.
pub struct RowIter<'a> {
  row : &'a rows::Row<'a>,
  index : usize
}

impl <'a> RowIter<'a> {
  /// Create a new row iterator starting at the beginning of the provided row
  pub fn new(row : &'a rows::Row) -> Self {
    RowIter {
      row   : row,
      index : 0
    }
  }
  /// Gets the next item in the row, using a `FromSql` instance to read it.
  /// If there is not a next item, returns `None`
  pub fn next<T>(&mut self) -> Option<T> where T : FromSql {
    let idx = self.index;
    self.index += 1;
    self.row.get(idx)
  }
}

/// Object representing a postgres-backed fact database instance
pub struct PgDB {
  conn              : Connection,
  pred_by_name      : HashMap<String, Predicate>,
  insert_by_name    : HashMap<String, String>,
  named_types       : HashMap<String, Type>
}

impl PgDB {
  /// Create a new PgDB object by passing in a Postgres connection string
  // TODO Add type parameters to call?
  // At the moment, persistence with custom types will result in failures
  // on a reconnect, so use a fresh database every time.
  // There's not a good way to persist custom types, so that fix will likely
  // come with optional parameters to seed types in at db startup.
  // TODO Should we be passing in a Connection object rather than a string?
  pub fn new(conn_str : &str) -> Result<PgDB> {
    let conn = try!(Connection::connect(conn_str, SslMode::None));

    // Create schemas
    try!(conn.execute("create schema if not exists facts", &[]));

    // Create Tables
    try!(conn.execute("create table if not exists predicates \
                       (pred_name varchar not null, \
                        ordinal int4 not null, \
                        type varchar not null)",
                      &[]));
    try!(conn.execute("create table if not exists rules \
                      (id serial, rule varchar not null)", &[]));

    // Create incremental PgDB object
    let mut db = PgDB {
      conn : conn,
      pred_by_name      : HashMap::new(),
      insert_by_name    : HashMap::new(),
      named_types       : types::default_types().iter().map(|type_| {
                            (type_.name().unwrap().to_owned(), type_.clone())
                          }).collect()
    };

    // Reload predicate cache
    {
      let pred_stmt = try!(db.conn.prepare(
        "select pred_name, type from predicates ORDER BY pred_name, ordinal"));
      let pred_types = try!(pred_stmt.query(&[]));
      for type_entry in pred_types.iter() {
        let name : String = type_entry.get(0);
        let h_type_str : String = type_entry.get(1);
        let h_type = match db.get_type(&h_type_str) {
          Some(ty) => ty,
          None => return Err(Box::new(
                               Error::Type(format!("Type not in registry: {}",
                                                   h_type_str))))
        };
        match db.pred_by_name.entry(name.clone()) {
          Vacant(entry) => {
            let mut types = Vec::new();
            types.push(h_type.clone());
            entry.insert(Predicate {
              name  : name.clone(),
              types : types
            });
          }
          Occupied(mut entry) => {
            entry.get_mut().types.push(h_type.clone());
          }
        }
      }
    }

    //Populate fact insert cache
    for pred in db.pred_by_name.clone().values() {
      &db.gen_insert_stmt(pred);
    }

    Ok(db)
  }

  // Generates a prebuilt insert statement for a given predicate, and stores
  // it in the cache so we don't have to rebuild it every time.
  // TODO: Is it possible for these to be stored prepared statements somehow?
  fn gen_insert_stmt(&mut self, pred : &Predicate) {
    let args : Vec<String> = pred.types.iter().enumerate().map(|(k,_)|{
      format!("${}", k + 1)
    }).collect();
    let stmt = format!("insert into facts.{} values ({})",
                       pred.name,
                       args.join(", "));
    self.insert_by_name.insert(pred.name.clone(), stmt);
  }

  // Persist a predicate into the database
  // This function is internal because it does not add it to the object, it
  // _only_ puts record of the predicate into the database.
  fn insert_predicate(&self, pred : &Predicate) -> Result<()> {
    let &Predicate {ref name, ref types} = pred;
    for (ordinal, type_) in types.iter().enumerate() {
      try!(self.conn.execute("insert into predicates \
                              (pred_name, type, ordinal) \
                              values ($1, $2, $3)",
                             &[name,
                               &type_.name().unwrap(),
                               &(ordinal as i32)]));
    }
    let table_str = types.iter().flat_map(|type_| {type_.repr()}).enumerate()
                          .map(|(ord, repr)| {
                            format!("arg{} {}", ord, repr)
                          }).collect::<Vec<_>>().join(", ");
    try!(self.conn.execute(
           &format!("create table facts.{} ({})", name, table_str),
           &[]));
    Ok(())
  }
}
impl FactDB for PgDB {
  /// Adds a new fact to the database, returning false if the fact was already
  /// present in the database, and true if it was inserted.
  fn insert_fact(&mut self, fact : &Fact) -> Result<bool> {
    let stmt : String = try!(self.insert_by_name
      .get(&fact.pred_name)
      .ok_or(Error::Internal("Insert Statement Missing"
                           .to_string()))).clone();
    Ok(try!(self.conn.execute(&stmt,
                              &fact.args.iter().flat_map(|x|{
                                x.to_sql().into_iter()
                              }).collect::<Vec<_>>())) > 0)
  }

  /// Registers a new type with the database.
  /// This is unstable, and will likely need to be moved to the initialization
  /// of the database object in order to allow reconnecting to an existing
  /// database.
  fn add_type(&mut self, type_ : Type) -> Result<()> {
    let name = type_.name().unwrap();
    if !self.named_types.contains_key(name) {
      self.named_types.insert(name.to_owned(), type_.clone());
      Ok(())
    } else {
      Err(Box::new(Error::Type(format!("{} already registered", name))))
    }
  }

  /// Looks for a named type in the database's registry.
  /// This function is primarily useful for the DSL shorthand for constructing
  /// queries, since it allows you to use names of types when declaring
  /// functions rather than type objects.
  fn get_type(&self, type_str : &str) -> Option<Type> {
    self.named_types.get(type_str).map(|x|{x.clone()})
  }

  /// Fetches a predicate by name
  fn get_predicate(&self, pred_name : &str) -> Option<&Predicate> {
    self.pred_by_name.get(pred_name)
  }

  /// Persists a predicate by name
  /// The name *must* consist only of lower case ASCII and _, anything else
  /// will be rejected. This restriction is because the predicate name is
  /// currently used to construct the table name.
  ///
  /// In the future, this restriction could be lifted by generating table
  /// names rather than using the names of predicates, but this helps a lot
  /// with debugging for now.
  // TODO lift restriction on predicate names
  fn new_predicate(&mut self, pred : &Predicate) -> Result<()> {
    // The predicate name is used as a table name, check it for legality
    if !valid_name(&pred.name) {
      return Err(Box::new(Error::Arg(
              "Invalid name: Use lowercase and underscores only".to_string())))
    }
    // If this predicate was already registered, fail
    if self.pred_by_name.contains_key(&pred.name) {
      return Err(Box::new(Error::Arg(
              format!("Predicate {} already registered.", &pred.name))))
    }

    try!(self.insert_predicate(&pred));
    self.gen_insert_stmt(&pred);
    self.pred_by_name.insert(pred.name.clone(), pred.clone());
    Ok(())
  }

  /// Attempt to match the right hand side of a datalog rule against the
  /// database, returning a list of solution assignments to the bound
  /// variables.
  fn search_facts(&self, query : &Vec<Clause>)
    -> Result<Vec<Vec<Value>>> {
    // Check there is at least one clause
    if query.len() == 0 {
      return Err(Box::new(Error::Arg("Empty search query".to_string())));
    };

    // Check that clauses:
    // * Have sequential variables
    // * Reference predicates in the database
    // * Only unify variables of equal type
    {
      let mut var_type : Vec<Type> = Vec::new();
      for clause in query.iter() {
        let pred = match self.pred_by_name.get(&clause.pred_name) {
          Some(pred) => pred,
          None => return Err(Box::new(Error::Arg(
                  format!("{} is not a registered predicate.",
                          clause.pred_name)))),
        };
        for (idx, slot) in clause.args.iter().enumerate() {
          match *slot {
              MatchExpr::Unbound
            | MatchExpr::Const(_) => (),
              MatchExpr::Var(v) => {
                let v = v as usize;
                if v == var_type.len() {
                  var_type.push(pred.types[idx].clone())
                } else if v > var_type.len() {
                  return Err(Box::new(Error::Arg(format!("Hole between {} and {} in variable numbering.", var_type.len() - 1, v))))
                } else if var_type[v] != pred.types[idx].clone() {
                  return Err(Box::new(Error::Arg(format!("Variable {} attempt to unify incompatible types {:?} and {:?}", v, var_type[v], pred.types[idx]))))
                }
              }
          }
        }
      }
    }

    // Actually build and execute the query
    let mut tables = Vec::new();    // Predicate names involved in the query,
                                    // in the sequence they appear
    let mut restricts = Vec::new(); // Unification expressions, indexed by
                                    // which join they belong on.
    let mut var_names = Vec::new(); // Translation of variable numbers to
                                    // sql exprs
    let mut var_types = Vec::new(); // Translation of variable numbers to
                                    // Types
    let mut where_clause = Vec::new(); // Constant comparisons
    let mut vals : Vec<&ToSql> = Vec::new(); // Values to be quoted into the
                                             // prepared statement

    for (idxc, clause) in query.iter().enumerate() {
      // The clause refers to a table named by the predicate
      let table_name = format!("facts.{}", clause.pred_name);
      // We will refer to it by a numbered alias, to make joining easier
      let alias_name = format!("t{}", idxc);
      let mut clause_elements = Vec::new();
      for (idx, arg) in clause.args.iter().enumerate() {
        match arg {
          &MatchExpr::Unbound => (),
          &MatchExpr::Var(var) => if var >= var_names.len() {
              // This situation means it's the first occurrence of the variable
              // We record this definition as the canonical definition for use
              // in the select, and store the type to know how to extract it.
              var_names.push(
                format!("{}.arg{}", alias_name, idx));
              var_types.push(&self.pred_by_name[&clause.pred_name].types[idx]);
            } else {
              // The variable has occurred correctly, so we add it being equal
              // to the canonical definition to the join clause for this table
              let piece = format!("{}.arg{} = {}", alias_name, idx,
                                  var_names[var]);
              clause_elements.push(piece);
            },
          &MatchExpr::Const(ref val) => {
            // Since we're comparing against a constant, this restriction can
            // go in the where clause.
            // I stash the value in a buffer for later use with the prepared
            // statement, and put the index into the buffer into the where
            // clause chunk.
            vals.extend(val.to_sql());
            where_clause.push(
              format!("{}.arg{} = ${}", alias_name, idx, vals.len()));
          }
        }
      }
      restricts.push(clause_elements);
      tables.push(format!("{} as {}", table_name, alias_name));
    }
    // Make sure we're never empty on bound variables. If we are, we will get
    // SELECT FROM
    // which will not work.
    var_names.push("0".to_string());

    let vars = format!("{}", var_names.join(", "));
    tables.reverse();
    restricts.reverse();
    let main_table = tables.pop().unwrap();
    where_clause.append(&mut restricts.pop().unwrap());
    let join_query =
      tables.iter().zip(restricts.iter()).map(|(table, join)| {
        if join.len() == 0 {
          format!("JOIN {} ", table)
        } else {
          format!("JOIN {} ON {}", table, join.join(" AND "))
        }
      }).collect::<Vec<_>>().join(" ");
    let where_clause = {
      if where_clause.len() == 0 {
        String::new()
      } else {
        format!("WHERE {}", where_clause.join(" AND "))
      }
    };
    let raw_stmt =
      format!("SELECT {} FROM {} {} {}",
              vars, main_table, join_query,
              where_clause);
    println!("{}", raw_stmt);
    let stmt = try!(self.conn.prepare(&raw_stmt));
    let rows = try!(stmt.query(&vals));

    let mut anss : Vec<Vec<Value>> = rows.iter().map(|row| {
      let mut row_iter = RowIter::new(&row);
      var_types.iter().map(|type_| {
        type_.extract(&mut row_iter)
      }).collect()
    }).collect();

    // TODO: Understand why this is necessary, if it should be necessary.
    anss.dedup();
    Ok(anss)
  }
}

fn valid_name(name : &String) -> bool {
  name.chars().all( |ch| match ch { 'a'...'z' | '_' => true, _ => false } )
}
<commit_msg>Remove debugging print<commit_after>//! Postgres-based Fact Database
//!
//! # Design Notes
//!
//! ## Scope
//!
//! The general philsophy is that things having to do with persistence go here,
//! while things related to non-persistent components go in `holmes::engine`.
//!
//! In the long run, we would like to persist nearly everything in the
//! database, so that a client-server model can one bay restored. However,
//! in the short term this has little benefit, so only the items needing to
//! use SQL for efficient execution are being included.
//!
//! The biggest hurdle here is the persistence of code:
//!
//! * How do we store Types?
//! * How do we store bound functions?
//!
//! One possible long term answer is Cap'n' Proto `SturdyRef`s
//!
//! ## Other Databases
//!
//! For the moment, this is the only implementation, and there are no others
//! on the horizon, so this interface is not abstract.
//!
//! The only major hurdle to using another backend would be figuring out how
//! to make the `dyn` module abstract over databases.
use std::collections::hash_map::HashMap;
use std::collections::hash_map::Entry::{Occupied, Vacant};

use postgres::{rows, Connection, SslMode};
use postgres::types::{FromSql, ToSql};

use engine::types::{Fact, Predicate, MatchExpr, Clause};

mod error;
pub mod dyn;

pub use self::error::Error;
use self::dyn::types;
use self::dyn::{Type, Value};
use fact_db::{FactDB, Result};

/// An iterator over a `postgres::rows::Row`.
/// It does not implement the normal iter interface because it does not have
/// a set item type, but it implements a similar interface for ease of use.
pub struct RowIter<'a> {
  row : &'a rows::Row<'a>,
  index : usize
}

impl <'a> RowIter<'a> {
  /// Create a new row iterator starting at the beginning of the provided row
  pub fn new(row : &'a rows::Row) -> Self {
    RowIter {
      row   : row,
      index : 0
    }
  }
  /// Gets the next item in the row, using a `FromSql` instance to read it.
  /// If there is not a next item, returns `None`
  pub fn next<T>(&mut self) -> Option<T> where T : FromSql {
    let idx = self.index;
    self.index += 1;
    self.row.get(idx)
  }
}

/// Object representing a postgres-backed fact database instance
pub struct PgDB {
  conn              : Connection,
  pred_by_name      : HashMap<String, Predicate>,
  insert_by_name    : HashMap<String, String>,
  named_types       : HashMap<String, Type>
}

impl PgDB {
  /// Create a new PgDB object by passing in a Postgres connection string
  // TODO Add type parameters to call?
  // At the moment, persistence with custom types will result in failures
  // on a reconnect, so use a fresh database every time.
  // There's not a good way to persist custom types, so that fix will likely
  // come with optional parameters to seed types in at db startup.
  // TODO Should we be passing in a Connection object rather than a string?
  pub fn new(conn_str : &str) -> Result<PgDB> {
    let conn = try!(Connection::connect(conn_str, SslMode::None));

    // Create schemas
    try!(conn.execute("create schema if not exists facts", &[]));

    // Create Tables
    try!(conn.execute("create table if not exists predicates \
                       (pred_name varchar not null, \
                        ordinal int4 not null, \
                        type varchar not null)",
                      &[]));
    try!(conn.execute("create table if not exists rules \
                      (id serial, rule varchar not null)", &[]));

    // Create incremental PgDB object
    let mut db = PgDB {
      conn : conn,
      pred_by_name      : HashMap::new(),
      insert_by_name    : HashMap::new(),
      named_types       : types::default_types().iter().map(|type_| {
                            (type_.name().unwrap().to_owned(), type_.clone())
                          }).collect()
    };

    // Reload predicate cache
    {
      let pred_stmt = try!(db.conn.prepare(
        "select pred_name, type from predicates ORDER BY pred_name, ordinal"));
      let pred_types = try!(pred_stmt.query(&[]));
      for type_entry in pred_types.iter() {
        let name : String = type_entry.get(0);
        let h_type_str : String = type_entry.get(1);
        let h_type = match db.get_type(&h_type_str) {
          Some(ty) => ty,
          None => return Err(Box::new(
                               Error::Type(format!("Type not in registry: {}",
                                                   h_type_str))))
        };
        match db.pred_by_name.entry(name.clone()) {
          Vacant(entry) => {
            let mut types = Vec::new();
            types.push(h_type.clone());
            entry.insert(Predicate {
              name  : name.clone(),
              types : types
            });
          }
          Occupied(mut entry) => {
            entry.get_mut().types.push(h_type.clone());
          }
        }
      }
    }

    //Populate fact insert cache
    for pred in db.pred_by_name.clone().values() {
      &db.gen_insert_stmt(pred);
    }

    Ok(db)
  }

  // Generates a prebuilt insert statement for a given predicate, and stores
  // it in the cache so we don't have to rebuild it every time.
  // TODO: Is it possible for these to be stored prepared statements somehow?
  fn gen_insert_stmt(&mut self, pred : &Predicate) {
    let args : Vec<String> = pred.types.iter().enumerate().map(|(k,_)|{
      format!("${}", k + 1)
    }).collect();
    let stmt = format!("insert into facts.{} values ({})",
                       pred.name,
                       args.join(", "));
    self.insert_by_name.insert(pred.name.clone(), stmt);
  }

  // Persist a predicate into the database
  // This function is internal because it does not add it to the object, it
  // _only_ puts record of the predicate into the database.
  fn insert_predicate(&self, pred : &Predicate) -> Result<()> {
    let &Predicate {ref name, ref types} = pred;
    for (ordinal, type_) in types.iter().enumerate() {
      try!(self.conn.execute("insert into predicates \
                              (pred_name, type, ordinal) \
                              values ($1, $2, $3)",
                             &[name,
                               &type_.name().unwrap(),
                               &(ordinal as i32)]));
    }
    let table_str = types.iter().flat_map(|type_| {type_.repr()}).enumerate()
                          .map(|(ord, repr)| {
                            format!("arg{} {}", ord, repr)
                          }).collect::<Vec<_>>().join(", ");
    try!(self.conn.execute(
           &format!("create table facts.{} ({})", name, table_str),
           &[]));
    Ok(())
  }
}
impl FactDB for PgDB {
  /// Adds a new fact to the database, returning false if the fact was already
  /// present in the database, and true if it was inserted.
  fn insert_fact(&mut self, fact : &Fact) -> Result<bool> {
    let stmt : String = try!(self.insert_by_name
      .get(&fact.pred_name)
      .ok_or(Error::Internal("Insert Statement Missing"
                           .to_string()))).clone();
    Ok(try!(self.conn.execute(&stmt,
                              &fact.args.iter().flat_map(|x|{
                                x.to_sql().into_iter()
                              }).collect::<Vec<_>>())) > 0)
  }

  /// Registers a new type with the database.
  /// This is unstable, and will likely need to be moved to the initialization
  /// of the database object in order to allow reconnecting to an existing
  /// database.
  fn add_type(&mut self, type_ : Type) -> Result<()> {
    let name = type_.name().unwrap();
    if !self.named_types.contains_key(name) {
      self.named_types.insert(name.to_owned(), type_.clone());
      Ok(())
    } else {
      Err(Box::new(Error::Type(format!("{} already registered", name))))
    }
  }

  /// Looks for a named type in the database's registry.
  /// This function is primarily useful for the DSL shorthand for constructing
  /// queries, since it allows you to use names of types when declaring
  /// functions rather than type objects.
  fn get_type(&self, type_str : &str) -> Option<Type> {
    self.named_types.get(type_str).map(|x|{x.clone()})
  }

  /// Fetches a predicate by name
  fn get_predicate(&self, pred_name : &str) -> Option<&Predicate> {
    self.pred_by_name.get(pred_name)
  }

  /// Persists a predicate by name
  /// The name *must* consist only of lower case ASCII and _, anything else
  /// will be rejected. This restriction is because the predicate name is
  /// currently used to construct the table name.
  ///
  /// In the future, this restriction could be lifted by generating table
  /// names rather than using the names of predicates, but this helps a lot
  /// with debugging for now.
  // TODO lift restriction on predicate names
  fn new_predicate(&mut self, pred : &Predicate) -> Result<()> {
    // The predicate name is used as a table name, check it for legality
    if !valid_name(&pred.name) {
      return Err(Box::new(Error::Arg(
              "Invalid name: Use lowercase and underscores only".to_string())))
    }
    // If this predicate was already registered, fail
    if self.pred_by_name.contains_key(&pred.name) {
      return Err(Box::new(Error::Arg(
              format!("Predicate {} already registered.", &pred.name))))
    }

    try!(self.insert_predicate(&pred));
    self.gen_insert_stmt(&pred);
    self.pred_by_name.insert(pred.name.clone(), pred.clone());
    Ok(())
  }

  /// Attempt to match the right hand side of a datalog rule against the
  /// database, returning a list of solution assignments to the bound
  /// variables.
  fn search_facts(&self, query : &Vec<Clause>)
    -> Result<Vec<Vec<Value>>> {
    // Check there is at least one clause
    if query.len() == 0 {
      return Err(Box::new(Error::Arg("Empty search query".to_string())));
    };

    // Check that clauses:
    // * Have sequential variables
    // * Reference predicates in the database
    // * Only unify variables of equal type
    {
      let mut var_type : Vec<Type> = Vec::new();
      for clause in query.iter() {
        let pred = match self.pred_by_name.get(&clause.pred_name) {
          Some(pred) => pred,
          None => return Err(Box::new(Error::Arg(
                  format!("{} is not a registered predicate.",
                          clause.pred_name)))),
        };
        for (idx, slot) in clause.args.iter().enumerate() {
          match *slot {
              MatchExpr::Unbound
            | MatchExpr::Const(_) => (),
              MatchExpr::Var(v) => {
                let v = v as usize;
                if v == var_type.len() {
                  var_type.push(pred.types[idx].clone())
                } else if v > var_type.len() {
                  return Err(Box::new(Error::Arg(format!("Hole between {} and {} in variable numbering.", var_type.len() - 1, v))))
                } else if var_type[v] != pred.types[idx].clone() {
                  return Err(Box::new(Error::Arg(format!("Variable {} attempt to unify incompatible types {:?} and {:?}", v, var_type[v], pred.types[idx]))))
                }
              }
          }
        }
      }
    }

    // Actually build and execute the query
    let mut tables = Vec::new();    // Predicate names involved in the query,
                                    // in the sequence they appear
    let mut restricts = Vec::new(); // Unification expressions, indexed by
                                    // which join they belong on.
    let mut var_names = Vec::new(); // Translation of variable numbers to
                                    // sql exprs
    let mut var_types = Vec::new(); // Translation of variable numbers to
                                    // Types
    let mut where_clause = Vec::new(); // Constant comparisons
    let mut vals : Vec<&ToSql> = Vec::new(); // Values to be quoted into the
                                             // prepared statement

    for (idxc, clause) in query.iter().enumerate() {
      // The clause refers to a table named by the predicate
      let table_name = format!("facts.{}", clause.pred_name);
      // We will refer to it by a numbered alias, to make joining easier
      let alias_name = format!("t{}", idxc);
      let mut clause_elements = Vec::new();
      for (idx, arg) in clause.args.iter().enumerate() {
        match arg {
          &MatchExpr::Unbound => (),
          &MatchExpr::Var(var) => if var >= var_names.len() {
              // This situation means it's the first occurrence of the variable
              // We record this definition as the canonical definition for use
              // in the select, and store the type to know how to extract it.
              var_names.push(
                format!("{}.arg{}", alias_name, idx));
              var_types.push(&self.pred_by_name[&clause.pred_name].types[idx]);
            } else {
              // The variable has occurred correctly, so we add it being equal
              // to the canonical definition to the join clause for this table
              let piece = format!("{}.arg{} = {}", alias_name, idx,
                                  var_names[var]);
              clause_elements.push(piece);
            },
          &MatchExpr::Const(ref val) => {
            // Since we're comparing against a constant, this restriction can
            // go in the where clause.
            // I stash the value in a buffer for later use with the prepared
            // statement, and put the index into the buffer into the where
            // clause chunk.
            vals.extend(val.to_sql());
            where_clause.push(
              format!("{}.arg{} = ${}", alias_name, idx, vals.len()));
          }
        }
      }
      restricts.push(clause_elements);
      tables.push(format!("{} as {}", table_name, alias_name));
    }
    // Make sure we're never empty on bound variables. If we are, we will get
    // SELECT FROM
    // which will not work.
    var_names.push("0".to_string());

    let vars = format!("{}", var_names.join(", "));
    tables.reverse();
    restricts.reverse();
    let main_table = tables.pop().unwrap();
    where_clause.append(&mut restricts.pop().unwrap());
    let join_query =
      tables.iter().zip(restricts.iter()).map(|(table, join)| {
        if join.len() == 0 {
          format!("JOIN {} ", table)
        } else {
          format!("JOIN {} ON {}", table, join.join(" AND "))
        }
      }).collect::<Vec<_>>().join(" ");
    let where_clause = {
      if where_clause.len() == 0 {
        String::new()
      } else {
        format!("WHERE {}", where_clause.join(" AND "))
      }
    };
    let raw_stmt =
      format!("SELECT {} FROM {} {} {}",
              vars, main_table, join_query,
              where_clause);
    let stmt = try!(self.conn.prepare(&raw_stmt));
    let rows = try!(stmt.query(&vals));

    let mut anss : Vec<Vec<Value>> = rows.iter().map(|row| {
      let mut row_iter = RowIter::new(&row);
      var_types.iter().map(|type_| {
        type_.extract(&mut row_iter)
      }).collect()
    }).collect();

    // TODO: Understand why this is necessary, if it should be necessary.
    anss.dedup();
    Ok(anss)
  }
}

fn valid_name(name : &String) -> bool {
  name.chars().all( |ch| match ch { 'a'...'z' | '_' => true, _ => false } )
}
<|endoftext|>


<commit_before>use core::fmt;

/// x86 Exception description (see also Intel Vol. 3a Chapter 6).
#[derive(Debug)]
pub struct InterruptDescription {
    pub vector: u8,
    pub mnemonic: &'static str,
    pub description: &'static str,
    pub irqtype: &'static str,
    pub source: &'static str
}

impl fmt::Display for InterruptDescription {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{} ({}, vec={}) {} {}", self.mnemonic, self.irqtype, self.vector, self.description, self.source)
    }
}

/// x86 External Interrupts (1-16).
pub static EXCEPTIONS: [InterruptDescription; 15] = [
    InterruptDescription{ vector: 0,  mnemonic: "#DE", description: "Divide Error", irqtype: "Fault", source: "DIV and IDIV instructions." },
    InterruptDescription{ vector: 1,  mnemonic: "#DB", description: "RESERVED", irqtype: "Fault/ Trap", source: "For Intel use only." },
    InterruptDescription{ vector: 2,  mnemonic: "NMI", description: "Interrupt", irqtype: "Interrupt", source: "Nonmaskable external interrupt." },
    InterruptDescription{ vector: 3,  mnemonic: "#BP", description: "Breakpoint", irqtype: "Trap", source: "INT 3 instruction." },
    InterruptDescription{ vector: 4,  mnemonic: "#OF", description: "Overflow", irqtype: "Trap", source: "INTO instruction." },
    InterruptDescription{ vector: 5,  mnemonic: "#BR", description: "BOUND Range Exceeded", irqtype: "Fault", source: "BOUND instruction." },
    InterruptDescription{ vector: 6,  mnemonic: "#UD", description: "Invalid Opcode (Undefined Opcode)", irqtype: "Fault", source: "UD2 instruction or reserved opcode." },
    InterruptDescription{ vector: 7,  mnemonic: "#NM", description: "Device Not Available (No Math Coprocessor)", irqtype: "Fault", source: "Floating-point or WAIT/FWAIT instruction." },
    InterruptDescription{ vector: 8,  mnemonic: "#DF", description: "Double Fault", irqtype: "Abort", source: "Any instruction that can generate an exception, an NMI, or an INTR." },
    InterruptDescription{ vector: 9,  mnemonic: ""   , description: "Coprocessor Segment Overrun", irqtype: "Fault", source: "Floating-point instruction." },
    InterruptDescription{ vector: 10, mnemonic: "#TS", description: "Invalid TSS", irqtype: "Fault", source: "Task switch or TSS access." },
    InterruptDescription{ vector: 11, mnemonic: "#NP", description: "Segment Not Present", irqtype: "Fault", source: "Loading segment registers or accessing system segments." },
    InterruptDescription{ vector: 12, mnemonic: "#SS", description: "Stack-Segment Fault", irqtype: "Fault", source: "Stack operations and SS register loads." },
    InterruptDescription{ vector: 13, mnemonic: "#GP", description: "General Protection", irqtype: "Fault", source: "Any memory reference and other protection checks." },
    InterruptDescription{ vector: 14, mnemonic: "#PF", description: "Page Fault", irqtype: "Fault", source: "Any memory reference." }
];

/// Enable Interrupts.
pub unsafe fn enable()  {
    asm!("sti");
}

/// Disable Interrupts.
pub unsafe fn disable()  {
    asm!("cli");
}

/// Generate a software interrupt.
/// This is a macro argument needs to be an immediate.
#[macro_export]
macro_rules! int {
    ( $x:expr ) => {
        {
            asm!("int $0" :: "N" ($x));
        }
    };
}

/// A struct describing an interrupt gate.
#[derive(Debug, Clone)]
#[repr(packed)]
pub struct IdtEntry {
    /// Lower 16 bits of ISR.
    pub base_lo: u16,
    /// Segment selector.
    pub sel: u16,
    /// This must always be zero.
    pub res0: u8,
    /// Flags.
    pub flags: u8,
    /// The upper 48 bits of ISR (the last 16 bits must be zero).
    pub base_hi: u64,
    /// Must be zero.
    pub res1: u16
}

/// A struct describing a pointer to an array of interrupt handlers.
/// This is in a format suitable for giving to 'lidt'.
#[derive(Debug)]
#[repr(packed)]
pub struct IdtPointer {
   /// Size of the IDT.
   pub limit: u16,
   /// Pointer to the memory region containing the IDT.
   pub base: u64
}

/// Load IDT table.
pub unsafe fn lidt(idt: &IdtPointer) {
    asm!("lidt ($0)" :: "r" (idt));
}

<commit_msg>Added copy trait for IdtEntry.<commit_after>use core::fmt;

/// x86 Exception description (see also Intel Vol. 3a Chapter 6).
#[derive(Debug)]
pub struct InterruptDescription {
    pub vector: u8,
    pub mnemonic: &'static str,
    pub description: &'static str,
    pub irqtype: &'static str,
    pub source: &'static str
}

impl fmt::Display for InterruptDescription {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{} ({}, vec={}) {} {}", self.mnemonic, self.irqtype, self.vector, self.description, self.source)
    }
}

/// x86 External Interrupts (1-16).
pub static EXCEPTIONS: [InterruptDescription; 15] = [
    InterruptDescription{ vector: 0,  mnemonic: "#DE", description: "Divide Error", irqtype: "Fault", source: "DIV and IDIV instructions." },
    InterruptDescription{ vector: 1,  mnemonic: "#DB", description: "RESERVED", irqtype: "Fault/ Trap", source: "For Intel use only." },
    InterruptDescription{ vector: 2,  mnemonic: "NMI", description: "Interrupt", irqtype: "Interrupt", source: "Nonmaskable external interrupt." },
    InterruptDescription{ vector: 3,  mnemonic: "#BP", description: "Breakpoint", irqtype: "Trap", source: "INT 3 instruction." },
    InterruptDescription{ vector: 4,  mnemonic: "#OF", description: "Overflow", irqtype: "Trap", source: "INTO instruction." },
    InterruptDescription{ vector: 5,  mnemonic: "#BR", description: "BOUND Range Exceeded", irqtype: "Fault", source: "BOUND instruction." },
    InterruptDescription{ vector: 6,  mnemonic: "#UD", description: "Invalid Opcode (Undefined Opcode)", irqtype: "Fault", source: "UD2 instruction or reserved opcode." },
    InterruptDescription{ vector: 7,  mnemonic: "#NM", description: "Device Not Available (No Math Coprocessor)", irqtype: "Fault", source: "Floating-point or WAIT/FWAIT instruction." },
    InterruptDescription{ vector: 8,  mnemonic: "#DF", description: "Double Fault", irqtype: "Abort", source: "Any instruction that can generate an exception, an NMI, or an INTR." },
    InterruptDescription{ vector: 9,  mnemonic: ""   , description: "Coprocessor Segment Overrun", irqtype: "Fault", source: "Floating-point instruction." },
    InterruptDescription{ vector: 10, mnemonic: "#TS", description: "Invalid TSS", irqtype: "Fault", source: "Task switch or TSS access." },
    InterruptDescription{ vector: 11, mnemonic: "#NP", description: "Segment Not Present", irqtype: "Fault", source: "Loading segment registers or accessing system segments." },
    InterruptDescription{ vector: 12, mnemonic: "#SS", description: "Stack-Segment Fault", irqtype: "Fault", source: "Stack operations and SS register loads." },
    InterruptDescription{ vector: 13, mnemonic: "#GP", description: "General Protection", irqtype: "Fault", source: "Any memory reference and other protection checks." },
    InterruptDescription{ vector: 14, mnemonic: "#PF", description: "Page Fault", irqtype: "Fault", source: "Any memory reference." }
];

/// Enable Interrupts.
pub unsafe fn enable()  {
    asm!("sti");
}

/// Disable Interrupts.
pub unsafe fn disable()  {
    asm!("cli");
}

/// Generate a software interrupt.
/// This is a macro argument needs to be an immediate.
#[macro_export]
macro_rules! int {
    ( $x:expr ) => {
        {
            asm!("int $0" :: "N" ($x));
        }
    };
}

/// A struct describing an interrupt gate.
#[derive(Debug, Clone, Copy)]
#[repr(packed)]
pub struct IdtEntry {
    /// Lower 16 bits of ISR.
    pub base_lo: u16,
    /// Segment selector.
    pub sel: u16,
    /// This must always be zero.
    pub res0: u8,
    /// Flags.
    pub flags: u8,
    /// The upper 48 bits of ISR (the last 16 bits must be zero).
    pub base_hi: u64,
    /// Must be zero.
    pub res1: u16
}

/// A struct describing a pointer to an array of interrupt handlers.
/// This is in a format suitable for giving to 'lidt'.
#[derive(Debug)]
#[repr(packed)]
pub struct IdtPointer {
   /// Size of the IDT.
   pub limit: u16,
   /// Pointer to the memory region containing the IDT.
   pub base: u64
}

/// Load IDT table.
pub unsafe fn lidt(idt: &IdtPointer) {
    asm!("lidt ($0)" :: "r" (idt));
}

<|endoftext|>


<commit_before>use addressable::Addressable;
use addressable::AddressableUtil;

#[cfg(test)]
use testmemory::TestMemory;

enum Operand {
	Implied,							// OPC				// operand implied
	Immediate(u8),						// OPC #$BB			// operand is byte (BB)
	Accumulator,						// OPC A			// operand is AC
	Relative(i8),						// OPC $BB			// branch target is PC + offset (BB), bit 7 signifies negative offset
	Absolute(u16),						// OPC $HHLL		// operand is address $HHLL
	AbsoluteIndexedWithX(u16),			// OPC $HHLL,X		// operand is address $HHLL incremented by X with carry
	AbsoluteIndexedWithY(u16),			// OPC $HHLL,Y		// operand is address $HHLL incremented by Y with carry
	Indirect(u16),						// OPC ($HHLL)		// operand is effective address; effective address is value of address
	ZeroPage(u8),						// OPC $LL			// operand is address $00LL
	ZeroPageIndexedWithX(u8),			// OPC $LL,X		// operand is address $00LL incremented by X, no page transition
	ZeroPageIndexedWithY(u8),			// OPC $LL,Y		// operand is address $00LL incremented by Y, no page transition
	ZeroPageIndexedWithXIndirect(u8),	// OPC ($LL,X)		// operand is effective address; effective address is $00LL incremented by X without carry
	ZeroPageIndirectIndexedWithY(u8),	// OPC ($LL),Y		// operand is effective address incremented by Y with carry; effective address is word at $00LL
}

impl Operand {
	fn indirect<M: Addressable<u16>> (mem: &M, addr: u16) -> u16 {
		// Get a little endian u16 from the given address, simulating the 6502 MSB bug
		// (i.e. take values from $c0ff/$c000 instead of $c0ff/$c100)
		let addr1 = (addr & 0xff00) | (addr+1 & 0x00ff);
		mem.get(addr) as u16 | mem.get(addr1) as u16 << 8
	}

	fn addr<M: Addressable<u16>> (&self, cpu: &Mos6502, mem: &M) -> u16 {
		match *self {
			Implied								=> fail!("mos65xx: Implied operand is never targetted to an address"),
			Immediate(_)						=> fail!("mos65xx: Immediate operand is never targetted to an address"),
			Accumulator							=> fail!("mos65xx: Accumulator operand is never targetted to an address"),
			Relative(offset)					=> cpu.pc + offset as u16,
			Absolute(addr)						=> addr,
			AbsoluteIndexedWithX(addr)			=> addr + cpu.x as u16,
			AbsoluteIndexedWithY(addr)			=> addr + cpu.y as u16,
			Indirect(addr)						=> Operand::indirect(mem, addr),
			ZeroPage(addr)						=> addr as u16,
			ZeroPageIndexedWithX(addr)			=> (addr + cpu.x) as u16,						// no page transition
			ZeroPageIndexedWithY(addr)			=> (addr + cpu.y) as u16,						// no page transition
			ZeroPageIndexedWithXIndirect(addr)	=> mem.get_le((addr + cpu.x) as u16),			// no page transition
			ZeroPageIndirectIndexedWithY(addr)	=> { let iaddr: u16 = mem.get_le(addr as u16); iaddr + cpu.y as u16 },
		}
	}

	fn get<M: Addressable<u16>> (&self, cpu: &Mos6502, mem: &M) -> u8 {
		match *self {
			Implied								=> fail!("mos65xx: Implied operand never has a value"),
			Immediate(val)						=> val,
			Accumulator							=> cpu.ac,
			Relative(_target)					=> fail!("mos65xx: Relative operand never has a value"),
			op									=> { let addr = op.addr(cpu, mem); mem.get(addr) },
		}
	}

	fn set<M: Addressable<u16>> (&self, cpu: &mut Mos6502, mem: &mut M, val: u8) {
		match *self {
			Implied								=> fail!("mos65xx: Implied operand never sets a value"),
			Immediate(_)						=> fail!("mos65xx: Immediate operand never sets a value"),
			Accumulator							=> cpu.ac = val,
			Relative(_target)					=> fail!("mos65xx: Relative operand never sets a value"),
			op									=> { let addr = op.addr(cpu, mem); mem.set(addr, val); },
		}
	}
}

// General information on 65xx: http://en.wikipedia.org/wiki/MOS_Technology_6510
// Web simulator and much info: http://e-tradition.net/bytes/6502/
// Good reference and overview: http://www.obelisk.demon.co.uk/index.html
// Processor bugs and caveats : http://www.textfiles.com/apple/6502.bugs.txt

static NMI_VECTOR: u16 = 0xfffa;
static RESET_VECTOR: u16 = 0xfffc;
static IRQ_VECTOR: u16 = 0xfffe;

pub struct Mos6502 {
	priv pc: u16,						// program counter
	priv ac: u8,						// accumulator
	priv x: u8,							// x register
	priv y: u8,							// y register
	priv sr: u8,						// status register (NV-BDIZC: Negative, oVerflow, Break, Decimal, Interrupt, Zero, Carry)
	priv sp: u8,						// stack pointer
}

impl Mos6502 {
	pub fn new () -> Mos6502 {
		Mos6502 { pc: 0, ac: 0, x: 0, y: 0, sr: 0, sp: 0 }
	}

	fn get_opcode<M: Addressable<u16>> (&self, mem: &M) -> u8 {
		mem.get(self.pc)
	}

	fn get_argument<M: Addressable<u16>, T: Int> (&self, mem: &M) -> T {
		mem.get_le(self.pc + 1)
	}

	pub fn reset<M: Addressable<u16>> (&mut self, mem: &M) {
		// On reset, the interrupt-disable flag is set (and the decimal flag is cleared in the CMOS version 65c02).
		// The other bits and all registers (including the stack pointer are unspecified and might contain random values.
		// Execution begins at the address pointed to by the reset vector at address $FFFC.
		debug!("mos65xx: Reset");
		self.pc = mem.get_le(RESET_VECTOR);
		self.sr = 0x24;
	}

	pub fn step<M: Addressable<u16>> (&mut self, mem: &mut M) -> uint {
		// TODO
		0
	}
}

pub struct Mos6510 {
	priv cpu: Mos6502,						// Core CPU
	priv port_ddr: u8,						// CPU port data direction register
	priv port_dat: u8,						// CPU port data register
}

impl Mos6510 {
	pub fn new () -> Mos6510 {
		// TODO: addresses $0000 (data direction) and $0001 (data) are hardwired for the processor I/O port
		Mos6510 { cpu: Mos6502::new(), port_ddr: 0, port_dat: 0 }
	}

	pub fn reset<M: Addressable<u16>> (&mut self, mem: &M) {
		self.cpu.reset(mem);
	}

	pub fn step<M: Addressable<u16>> (&mut self, mem: &mut M) -> uint {
		self.cpu.step(mem)
	}
}

#[test]
fn test_addressing_modes () {
	let mut cpu = Mos6502 { pc: 0x1337, ac: 0x88, x: 0x11, y: 0x22, sr: 0, sp: 0 };
	let mut mem = TestMemory::new::<u16>();
	// Immediate
	assert_eq!(Immediate(0x55).get(&cpu, &mem), 0x55);
	// Accumulator
	assert_eq!(Accumulator.get(&cpu, &mem), 0x88);
	Accumulator.set(&mut cpu, &mut mem, 0x99); assert_eq!(cpu.ac, 0x99);
	// Relative
	assert_eq!(Relative(0x33).addr(&cpu, &mem), 0x136a);
	assert_eq!(Relative(0x99).addr(&cpu, &mem), 0x12d0);
	// Absolute
	assert_eq!(Absolute(0x0123).addr(&cpu, &mem), 0x0123);
	assert_eq!(Absolute(0x0123).get(&cpu, &mem), 0x24);
	Absolute(0x0123).set(&mut cpu, &mut mem, 0x24);
	// AbsoluteIndexedWithX
	assert_eq!(AbsoluteIndexedWithX(0x0123).addr(&cpu, &mem), 0x0134);
	assert_eq!(AbsoluteIndexedWithX(0x0123).get(&cpu, &mem), 0x35);
	AbsoluteIndexedWithX(0x0123).set(&mut cpu, &mut mem, 0x35);
	// AbsoluteIndexedWithY
	assert_eq!(AbsoluteIndexedWithY(0x0123).addr(&cpu, &mem), 0x0145);
	assert_eq!(AbsoluteIndexedWithY(0x0123).get(&cpu, &mem), 0x46);
	AbsoluteIndexedWithY(0x0123).set(&mut cpu, &mut mem, 0x46);
	// Indirect
	assert_eq!(Indirect(0x0123).addr(&cpu, &mem), 0x2524);
	assert_eq!(Indirect(0x0123).get(&cpu, &mem), 0x49);
	Indirect(0x0123).set(&mut cpu, &mut mem, 0x49);
	// ZeroPage
	assert_eq!(ZeroPage(0x12).addr(&cpu, &mem), 0x0012);
	assert_eq!(ZeroPage(0x12).get(&cpu, &mem), 0x12);
	ZeroPage(0x12).set(&mut cpu, &mut mem, 0x12);
	// ZeroPageIndexedWithX
	assert_eq!(ZeroPageIndexedWithX(0x12).addr(&cpu, &mem), 0x0023);
	assert_eq!(ZeroPageIndexedWithX(0x12).get(&cpu, &mem), 0x23);
	ZeroPageIndexedWithX(0x12).set(&mut cpu, &mut mem, 0x23);
	// ZeroPageIndexedWithY
	assert_eq!(ZeroPageIndexedWithY(0x12).addr(&cpu, &mem), 0x0034);
	assert_eq!(ZeroPageIndexedWithY(0x12).get(&cpu, &mem), 0x34);
	ZeroPageIndexedWithY(0x12).set(&mut cpu, &mut mem, 0x34);
	// ZeroPageIndexedWithXIndirect
	assert_eq!(ZeroPageIndexedWithXIndirect(0x12).addr(&cpu, &mem), 0x2423);
	assert_eq!(ZeroPageIndexedWithXIndirect(0x12).get(&cpu, &mem), 0x47);
	ZeroPageIndexedWithXIndirect(0x12).set(&mut cpu, &mut mem, 0x47);
	// ZeroPageIndirectIndexedWithY
	assert_eq!(ZeroPageIndirectIndexedWithY(0x12).addr(&cpu, &mem), 0x1334);
	assert_eq!(ZeroPageIndirectIndexedWithY(0x12).get(&cpu, &mem), 0x47);
	ZeroPageIndirectIndexedWithY(0x12).set(&mut cpu, &mut mem, 0x47);
}

#[test]
fn test_indirect_addressing_bug () {
	let cpu = Mos6502 { pc: 0x1337, ac: 0x88, x: 0x11, y: 0x22, sr: 0, sp: 0 };
	let mem = TestMemory::new::<u16>();
	// Indirect jump to ($c0ff) will get erroneously address from $c0ff/$c000 instead of $c0ff/$c100
	assert_eq!(Indirect(0xc0ff).addr(&cpu, &mem), 0xc0bf);	// must be $c0bf instead of $c1bf
}
<commit_msg>Implemented CPU flags and helpers<commit_after>use addressable::Addressable;
use addressable::AddressableUtil;

#[cfg(test)]
use testmemory::TestMemory;

enum Operand {
	Implied,							// OPC				// operand implied
	Immediate(u8),						// OPC #$BB			// operand is byte (BB)
	Accumulator,						// OPC A			// operand is AC
	Relative(i8),						// OPC $BB			// branch target is PC + offset (BB), bit 7 signifies negative offset
	Absolute(u16),						// OPC $HHLL		// operand is address $HHLL
	AbsoluteIndexedWithX(u16),			// OPC $HHLL,X		// operand is address $HHLL incremented by X with carry
	AbsoluteIndexedWithY(u16),			// OPC $HHLL,Y		// operand is address $HHLL incremented by Y with carry
	Indirect(u16),						// OPC ($HHLL)		// operand is effective address; effective address is value of address
	ZeroPage(u8),						// OPC $LL			// operand is address $00LL
	ZeroPageIndexedWithX(u8),			// OPC $LL,X		// operand is address $00LL incremented by X, no page transition
	ZeroPageIndexedWithY(u8),			// OPC $LL,Y		// operand is address $00LL incremented by Y, no page transition
	ZeroPageIndexedWithXIndirect(u8),	// OPC ($LL,X)		// operand is effective address; effective address is $00LL incremented by X without carry
	ZeroPageIndirectIndexedWithY(u8),	// OPC ($LL),Y		// operand is effective address incremented by Y with carry; effective address is word at $00LL
}

impl Operand {
	fn indirect<M: Addressable<u16>> (mem: &M, addr: u16) -> u16 {
		// Get a little endian u16 from the given address, simulating the 6502 MSB bug
		// (i.e. take values from $c0ff/$c000 instead of $c0ff/$c100)
		let addr1 = (addr & 0xff00) | (addr+1 & 0x00ff);
		mem.get(addr) as u16 | mem.get(addr1) as u16 << 8
	}

	fn addr<M: Addressable<u16>> (&self, cpu: &Mos6502, mem: &M) -> u16 {
		match *self {
			Implied								=> fail!("mos65xx: Implied operand is never targetted to an address"),
			Immediate(_)						=> fail!("mos65xx: Immediate operand is never targetted to an address"),
			Accumulator							=> fail!("mos65xx: Accumulator operand is never targetted to an address"),
			Relative(offset)					=> cpu.pc + offset as u16,
			Absolute(addr)						=> addr,
			AbsoluteIndexedWithX(addr)			=> addr + cpu.x as u16,
			AbsoluteIndexedWithY(addr)			=> addr + cpu.y as u16,
			Indirect(addr)						=> Operand::indirect(mem, addr),
			ZeroPage(addr)						=> addr as u16,
			ZeroPageIndexedWithX(addr)			=> (addr + cpu.x) as u16,						// no page transition
			ZeroPageIndexedWithY(addr)			=> (addr + cpu.y) as u16,						// no page transition
			ZeroPageIndexedWithXIndirect(addr)	=> mem.get_le((addr + cpu.x) as u16),			// no page transition
			ZeroPageIndirectIndexedWithY(addr)	=> { let iaddr: u16 = mem.get_le(addr as u16); iaddr + cpu.y as u16 },
		}
	}

	fn get<M: Addressable<u16>> (&self, cpu: &Mos6502, mem: &M) -> u8 {
		match *self {
			Implied								=> fail!("mos65xx: Implied operand never has a value"),
			Immediate(val)						=> val,
			Accumulator							=> cpu.ac,
			Relative(_target)					=> fail!("mos65xx: Relative operand never has a value"),
			op									=> { let addr = op.addr(cpu, mem); mem.get(addr) },
		}
	}

	fn set<M: Addressable<u16>> (&self, cpu: &mut Mos6502, mem: &mut M, val: u8) {
		match *self {
			Implied								=> fail!("mos65xx: Implied operand never sets a value"),
			Immediate(_)						=> fail!("mos65xx: Immediate operand never sets a value"),
			Accumulator							=> cpu.ac = val,
			Relative(_target)					=> fail!("mos65xx: Relative operand never sets a value"),
			op									=> { let addr = op.addr(cpu, mem); mem.set(addr, val); },
		}
	}
}

// General information on 65xx: http://en.wikipedia.org/wiki/MOS_Technology_6510
// Web simulator and much info: http://e-tradition.net/bytes/6502/
// Good reference and overview: http://www.obelisk.demon.co.uk/index.html
// Processor bugs and caveats : http://www.textfiles.com/apple/6502.bugs.txt

enum StatusFlag {
	CarryFlag = 0,
	ZeroFlag = 1,
	InterruptDisableFlag = 2,
	DecimalFlag = 3,
	BreakFlag = 4,
	OverflowFlag = 6,
	NegativeFlag = 7,
}

static NMI_VECTOR:       u16 = 0xfffa;
static RESET_VECTOR:     u16 = 0xfffc;
static IRQ_VECTOR:       u16 = 0xfffe;

pub struct Mos6502 {
	priv pc: u16,						// program counter
	priv ac: u8,						// accumulator
	priv x: u8,							// x register
	priv y: u8,							// y register
	priv sr: u8,						// status register (NV-BDIZC: Negative, oVerflow, 1, Break, Decimal, Interrupt, Zero, Carry)
	priv sp: u8,						// stack pointer
}

impl Mos6502 {
	pub fn new () -> Mos6502 {
		Mos6502 { pc: 0, ac: 0, x: 0, y: 0, sr: 0, sp: 0 }
	}

	fn get_opcode<M: Addressable<u16>> (&self, mem: &M) -> u8 {
		mem.get(self.pc)
	}

	fn get_argument<M: Addressable<u16>, T: Int> (&self, mem: &M) -> T {
		mem.get_le(self.pc + 1)
	}

	fn get_flag (&self, flag: StatusFlag) -> bool {
		(self.sr & (1 << flag as u8)) != 0
	}

	fn set_flag (&mut self, flag: StatusFlag, on: bool) {
		if on {
			self.sr |= (1 << flag as u8);
		} else {
			self.sr &= !(1 << flag as u8);
		}
	}

	fn set_zn (&mut self, value: u8) -> u8 {
		self.set_flag(ZeroFlag, value == 0);
		self.set_flag(NegativeFlag, (value as i8) < 0);
		value
	}

	pub fn reset<M: Addressable<u16>> (&mut self, mem: &M) {
		// On reset, the interrupt-disable flag is set (and the decimal flag is cleared in the CMOS version 65c02).
		// The other bits and all registers (including the stack pointer are unspecified and might contain random values.
		// Execution begins at the address pointed to by the reset vector at address $FFFC.
		debug!("mos65xx: Reset");
		self.pc = mem.get_le(RESET_VECTOR);
		self.sr = 0x24;
	}

	pub fn step<M: Addressable<u16>> (&mut self, mem: &mut M) -> uint {
		// TODO
		0
	}
}

pub struct Mos6510 {
	priv cpu: Mos6502,						// Core CPU
	priv port_ddr: u8,						// CPU port data direction register
	priv port_dat: u8,						// CPU port data register
}

impl Mos6510 {
	pub fn new () -> Mos6510 {
		// TODO: addresses $0000 (data direction) and $0001 (data) are hardwired for the processor I/O port
		Mos6510 { cpu: Mos6502::new(), port_ddr: 0, port_dat: 0 }
	}

	pub fn reset<M: Addressable<u16>> (&mut self, mem: &M) {
		self.cpu.reset(mem);
	}

	pub fn step<M: Addressable<u16>> (&mut self, mem: &mut M) -> uint {
		self.cpu.step(mem)
	}
}

#[test]
fn test_addressing_modes () {
	let mut cpu = Mos6502 { pc: 0x1337, ac: 0x88, x: 0x11, y: 0x22, sr: 0, sp: 0 };
	let mut mem = TestMemory::new::<u16>();
	// Immediate
	assert_eq!(Immediate(0x55).get(&cpu, &mem), 0x55);
	// Accumulator
	assert_eq!(Accumulator.get(&cpu, &mem), 0x88);
	Accumulator.set(&mut cpu, &mut mem, 0x99); assert_eq!(cpu.ac, 0x99);
	// Relative
	assert_eq!(Relative(0x33).addr(&cpu, &mem), 0x136a);
	assert_eq!(Relative(0x99).addr(&cpu, &mem), 0x12d0);
	// Absolute
	assert_eq!(Absolute(0x0123).addr(&cpu, &mem), 0x0123);
	assert_eq!(Absolute(0x0123).get(&cpu, &mem), 0x24);
	Absolute(0x0123).set(&mut cpu, &mut mem, 0x24);
	// AbsoluteIndexedWithX
	assert_eq!(AbsoluteIndexedWithX(0x0123).addr(&cpu, &mem), 0x0134);
	assert_eq!(AbsoluteIndexedWithX(0x0123).get(&cpu, &mem), 0x35);
	AbsoluteIndexedWithX(0x0123).set(&mut cpu, &mut mem, 0x35);
	// AbsoluteIndexedWithY
	assert_eq!(AbsoluteIndexedWithY(0x0123).addr(&cpu, &mem), 0x0145);
	assert_eq!(AbsoluteIndexedWithY(0x0123).get(&cpu, &mem), 0x46);
	AbsoluteIndexedWithY(0x0123).set(&mut cpu, &mut mem, 0x46);
	// Indirect
	assert_eq!(Indirect(0x0123).addr(&cpu, &mem), 0x2524);
	assert_eq!(Indirect(0x0123).get(&cpu, &mem), 0x49);
	Indirect(0x0123).set(&mut cpu, &mut mem, 0x49);
	// ZeroPage
	assert_eq!(ZeroPage(0x12).addr(&cpu, &mem), 0x0012);
	assert_eq!(ZeroPage(0x12).get(&cpu, &mem), 0x12);
	ZeroPage(0x12).set(&mut cpu, &mut mem, 0x12);
	// ZeroPageIndexedWithX
	assert_eq!(ZeroPageIndexedWithX(0x12).addr(&cpu, &mem), 0x0023);
	assert_eq!(ZeroPageIndexedWithX(0x12).get(&cpu, &mem), 0x23);
	ZeroPageIndexedWithX(0x12).set(&mut cpu, &mut mem, 0x23);
	// ZeroPageIndexedWithY
	assert_eq!(ZeroPageIndexedWithY(0x12).addr(&cpu, &mem), 0x0034);
	assert_eq!(ZeroPageIndexedWithY(0x12).get(&cpu, &mem), 0x34);
	ZeroPageIndexedWithY(0x12).set(&mut cpu, &mut mem, 0x34);
	// ZeroPageIndexedWithXIndirect
	assert_eq!(ZeroPageIndexedWithXIndirect(0x12).addr(&cpu, &mem), 0x2423);
	assert_eq!(ZeroPageIndexedWithXIndirect(0x12).get(&cpu, &mem), 0x47);
	ZeroPageIndexedWithXIndirect(0x12).set(&mut cpu, &mut mem, 0x47);
	// ZeroPageIndirectIndexedWithY
	assert_eq!(ZeroPageIndirectIndexedWithY(0x12).addr(&cpu, &mem), 0x1334);
	assert_eq!(ZeroPageIndirectIndexedWithY(0x12).get(&cpu, &mem), 0x47);
	ZeroPageIndirectIndexedWithY(0x12).set(&mut cpu, &mut mem, 0x47);
}

#[test]
fn test_indirect_addressing_bug () {
	let cpu = Mos6502 { pc: 0x1337, ac: 0x88, x: 0x11, y: 0x22, sr: 0, sp: 0 };
	let mem = TestMemory::new::<u16>();
	// Indirect jump to ($c0ff) will get erroneously address from $c0ff/$c000 instead of $c0ff/$c100
	assert_eq!(Indirect(0xc0ff).addr(&cpu, &mem), 0xc0bf);	// must be $c0bf instead of $c1bf
}
<|endoftext|>