<commit_before>extern crate rustc_serialize;

pub use self::Ast::*;

use lexer::Lexer;
use lexer::Token;
use self::rustc_serialize::json::{Json};

use std::iter::Peekable;

/// Parses a JMESPath expression into an AST
pub fn parse(expr: &str) -> Result<Ast, ParseError> {
    Parser::new(expr).parse()
}

/// Represents the abstract syntax tree of a JMESPath expression.
#[derive(Clone, PartialEq, Debug)]
pub enum Ast {
    Comparison(Comparator, Box<Ast>, Box<Ast>),
    CurrentNode,
    Expref(Box<Ast>),
    Flatten(Box<Ast>),
    Function(char, Vec<(Box<Ast>, Box<Ast>)>),
    Identifier(String),
    Index(i32),
    Literal(Box<Json>),
    MultiList(Vec<Box<Ast>>),
    MultiHash(Vec<Vec<(char, Box<Ast>)>>),
    ArrayProjection(Box<Ast>, Box<Ast>),
    ObjectProjection(Box<Ast>, Box<Ast>),
    Or(Box<Ast>, Box<Ast>),
    Slice(Option<i32>, Option<i32>, Option<i32>),
    Subexpr(Box<Ast>, Box<Ast>),
    WildcardIndex,
    WildcardValues,
}

/// Comparators (i.e., less than, greater than, etc.)
#[derive(Clone, PartialEq, Debug)]
pub enum Comparator { Eq, Lt, Le, Ne, Ge, Gt }

/// Encountered when an invalid JMESPath expression is parsed.
#[derive(Clone, PartialEq, Debug)]
pub struct ParseError {
    /// The error message.
    msg: String,
    /// The line number of the error.
    line: usize,
    /// The column of the error.
    col: usize,
}

/// JMESPath parser. Returns an Ast
pub struct Parser<'a> {
    /// Peekable token stream
    stream: Peekable<Lexer<'a>>,
    /// Expression being parsed
    expr: String,
    /// The current token
    token: Token,
    /// The current character offset in the expression
    pos: usize,
}

impl<'a> Parser<'a> {
    // Constructs a new lexer using the given expression string.
    pub fn new(expr: &'a str) -> Parser<'a> {
        let mut stream = Lexer::new(expr).peekable();
        let tok0 = stream.next().unwrap();
        Parser {
            stream: stream,
            expr: expr.to_string(),
            token: tok0,
            pos: 0,
        }
    }

    /// Parses the expression into result containing an AST or ParseError.
    pub fn parse(&mut self) -> Result<Ast, ParseError> {
        let result = self.expr(0);
        let token = self.stream.next();
        // After parsing the expression, we should have reached the end of the
        // stream.
        if result.is_err()
        || token.is_none()
        || token.unwrap() == Token::Eof {
            result
        } else {
            return self.err("Did not reach token stream EOF".as_ref());
        }
    }

    /// Ensures that the next token in the token stream is one of the pipe
    /// separated token named provided as the edible argument (e.g.,
    /// "Identifier|Eof").
    fn expect(&mut self, edible: &str) -> Result<Ast, ParseError> {
        self.advance();
        // Get the string name of the token.
        let token_name = self.token.token_to_string();
        if edible.contains(&token_name) {
            return Ok(CurrentNode);
        }
        let msg = format!("Expected one of the following tokens: {:?}", edible);
        self.err(&msg)
    }

    /// Advances the cursor position, skipping any whitespace encountered.
    #[inline]
    fn advance(&mut self) {
        self.pos += self.token.size();
        loop {
            let tok = self.stream.next();
            match tok {
                None => break,
                Some(Token::Whitespace) => self.pos += self.token.size(),
                _ => {
                    self.token = tok.unwrap();
                    break
                }
            }
        }
    }

    /// Main parse function of the Pratt parser that parses while RBP < LBP
    pub fn expr(&mut self, rbp: usize) -> Result<Ast, ParseError> {
        // Parse the NUD token.
        let mut left = match self.token.clone() {
            Token::At               => self.nud_at(),
            Token::Identifier(s, _) => self.nud_identifier(s),
            Token::Star             => self.nud_star(),
            Token::Lbracket         => self.nud_lbracket(),
            Token::Flatten          => self.nud_flatten(),
            // Token::Literal(_, _)    => self.nud_literal(),
            // Token::Ampersand        => self.nud_ampersand(),
            // Token::Lbrace           => self.nud_lbrace(),
            // Token::Filter           => self.nud_filter(),
            _ => { return self.err(&"Unexpected NUD token"); }
        };

        // Parse any LED tokens with a higher binding power.
        while rbp < self.token.lbp() {
            left = match self.token {
                Token::Dot      => self.led_dot(left.unwrap()),
                Token::Lbracket => self.led_lbracket(left.unwrap()),
                Token::Flatten  => self.led_flatten(left.unwrap()),
                Token::Or       => self.led_or(left.unwrap()),
                Token::Pipe     => self.led_pipe(left.unwrap()),
                _ => {
                    self.err(&"Unexpected LED token");
                    break;
                }
            };
        }

        left
    }

    /// Returns a formatted ParseError with the given message.
    fn err(&self, msg: &str) -> Result<Ast, ParseError> {
        // Find each new line and create a formatted error message.
        let mut line = 0;
        let mut col = self.pos;
        Err(ParseError {
            msg: format!("Error at {:?} token, {}: {}",
                         self.token, self.pos, msg),
            col: line,
            line: col
        })
    }

    /// Examples: "@"
    fn nud_at(&mut self) -> Result<Ast, ParseError> {
        self.advance();
        Ok(CurrentNode)
    }

    /// Examples: "Foo"
    fn nud_identifier(&mut self, s: String) -> Result<Ast, ParseError> {
        self.advance();
        Ok(Identifier(s))
    }

    /// Examples: "[0]", "[*]", "[a, b]", "[0:1]", etc...
    fn nud_lbracket(&mut self) -> Result<Ast, ParseError> {
        self.advance();
        match self.token {
            Token::Number(_, _) | Token::Colon => self.parse_array_index(),
            Token::Star => {
                if self.stream.peek() != Some(&Token::Rbracket) {
                    return self.parse_multi_list();
                }
                try!(self.expect("Star"));
                self.parse_wildcard_index()
            },
            _ => self.parse_multi_list()
        }
    }

    /// Examples: foo[*], foo[0], foo[:-1], etc.
    fn led_lbracket(&mut self, lhs: Ast) -> Result<Ast, ParseError> {
        try!(self.expect("Number|Colon|Star"));
        match self.token {
            Token::Number(_, _) | Token::Colon => self.parse_array_index(),
            _ => self.parse_wildcard_index()
        }
    }

    /// Examples: "*" (e.g., "* | *" would be a pipe containing two nud stars)
    fn nud_star(&mut self) -> Result<Ast, ParseError> {
        self.advance();
        self.parse_wildcard_values(CurrentNode)
    }

    /// Examples: "[]". Turns it into a LED flatten (i.e., "@[]").
    fn nud_flatten(&mut self) -> Result<Ast, ParseError> {
        self.led_flatten(CurrentNode)
    }

    /// Creates a Projection AST node for a flatten token.
    fn led_flatten(&mut self, lhs: Ast) -> Result<Ast, ParseError> {
        let rhs = try!(self.projection_rhs(Token::Flatten.lbp()));
        Ok(ArrayProjection(
            Box::new(Flatten(Box::new(lhs))),
            Box::new(rhs)
        ))
    }

    fn led_dot(&mut self, left: Ast) -> Result<Ast, ParseError> {
        let rhs = try!(self.parse_dot(Token::Dot.lbp()));
        Ok(Ast::Subexpr(Box::new(left), Box::new(rhs)))
    }

    fn led_or(&mut self, left: Ast) -> Result<Ast, ParseError> {
        self.advance();
        let rhs = try!(self.expr(Token::Or.lbp()));
        Ok(Or(Box::new(left), Box::new(rhs)))
    }

    fn led_pipe(&mut self, left: Ast) -> Result<Ast, ParseError> {
        self.advance();
        let rhs = try!(self.expr(Token::Pipe.lbp()));
        Ok(Subexpr(Box::new(left), Box::new(rhs)))
    }

    /// Parses the right hand side of a dot expression.
    fn parse_dot(&mut self, lbp: usize) -> Result<Ast, ParseError> {
        try!(self.expect("Identifier|Star|Lbrace|Lbracket|Ampersand|Filter"));
        match self.token {
            Token::Lbracket => { self.advance(); self.parse_multi_list() },
            _ => self.expr(lbp)
        }
    }

    /// Parses the right hand side of a projection, using the given LBP to
    /// determine when to stop consuming tokens.
    fn projection_rhs(&mut self, lbp: usize) -> Result<Ast, ParseError> {
        let lbp = self.token.lbp();
        match self.token {
            Token::Dot      => self.parse_dot(lbp),
            Token::Lbracket => self.expr(lbp),
            Token::Filter   => self.expr(lbp),
            _ if lbp < 10   => Ok(CurrentNode),
            _               => self.err("Syntax error found in projection")
        }
    }

    /// Creates a projection for "[*]"
    fn parse_wildcard_index(&mut self) -> Result<Ast, ParseError> {
        try!(self.expect("Rbracket"));
        let lhs = Box::new(CurrentNode);
        let rhs = try!(self.projection_rhs(Token::Star.lbp()));
        Ok(ArrayProjection(lhs, Box::new(rhs)))
    }

    /// Creates a projection for "*"
    fn parse_wildcard_values(&mut self, lhs: Ast) -> Result<Ast, ParseError> {
        let rhs = try!(self.projection_rhs(Token::Star.lbp()));
        Ok(ObjectProjection(Box::new(lhs), Box::new(rhs)))
    }

    /// Parses [0], [::-1], [0:-1], [0:1], etc...
    fn parse_array_index(&mut self) -> Result<Ast, ParseError> {
        let mut parts = [None, None, None];
        let mut pos = 0;
        loop {
            match self.token {
                Token::Colon => {
                    pos += 1;
                    if pos > 2 {
                        return self.err("Too many colons in slice expr");
                    }
                    try!(self.expect("Number|Colon|Rbracket"));
                },
                Token::Number(value, _) => {
                    parts[pos] = Some(value);
                    try!(self.expect("Colon|Rbracket"));
                },
                Token::Rbracket => { self.advance(); break; },
                _ => { return self.err("Unexpected token"); },
            }
        }

        if pos == 0 {
            // No colons were found, so this is a simple index extraction.
            Ok(Index(parts[0].unwrap()))
        } else {
            // Sliced array from start (e.g., [2:])
            let lhs = Slice(parts[0], parts[1], parts[2]);
            let rhs = try!(self.projection_rhs(Token::Star.lbp()));
            Ok(ArrayProjection(Box::new(lhs), Box::new(rhs)))
        }
    }

    /// Parses multi-select lists (e.g., "[foo, bar, baz]")
    fn parse_multi_list(&mut self) -> Result<Ast, ParseError> {
        let mut nodes = vec!();
        loop {
            nodes.push(Box::new(try!(self.expr(0))));
            match self.token {
                // Skip comma tokens
                Token::Comma => {
                    self.advance();
                    // The closing "Rbracket" token cannot follow a comma.
                    if self.token.token_to_string() == "Rbracket" {
                        return self.err("Unexpected token after comma");
                    }
                },
                // Terminal conditon
                Token::Rbracket => break,
                // Got to EOF without closing
                Token::Eof => return self.err("Unclosed multi-select"),
                _ => {}
            }
        }

        Ok(MultiList(nodes))
    }
}

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

    #[test] fn indentifier_test() {
        assert_eq!(parse("foo").unwrap(), Identifier("foo".to_string()));
    }

    #[test] fn current_node_test() {
        assert_eq!(parse("@").unwrap(), CurrentNode);
    }

    #[test] fn wildcard_values_test() {
        assert_eq!(parse("*").unwrap(),
                   ObjectProjection(Box::new(CurrentNode),
                                    Box::new(CurrentNode)));
    }

    #[test] fn dot_test() {
        assert!(parse("@.b").unwrap() == Subexpr(Box::new(CurrentNode),
                                                 Box::new(Identifier("b".to_string()))));
    }

    #[test] fn ensures_nud_token_is_valid_test() {

    }

    #[test] fn multi_list_test() {
        let l = MultiList(vec![Box::new(Identifier("a".to_string())),
                               Box::new(Identifier("b".to_string()))]);
        assert_eq!(parse("[a, b]").unwrap(), l);
    }

    #[test] fn multi_list_after_dot_test() {
        let l = MultiList(vec![Box::new(Identifier("a".to_string())),
                               Box::new(Identifier("b".to_string()))]);
        assert_eq!(parse("@.[a, b]").unwrap(), Subexpr(Box::new(CurrentNode), Box::new(l)));
    }

    #[test] fn parses_simple_index_extractions_test() {
        assert_eq!(parse("[0]").unwrap(), Index(0));
    }

    #[test] fn parses_single_element_slice_test() {
        assert_eq!(parse("[-1:]").unwrap(),
                   ArrayProjection(Box::new(Slice(Some(-1), None, None)),
                                   Box::new(CurrentNode)));
    }

    #[test] fn parses_double_element_slice_test() {
        assert_eq!(parse("[1:-1].a").unwrap(),
                   ArrayProjection(Box::new(Slice(Some(1), Some(-1), None)),
                                   Box::new(Identifier("a".to_string()))));
    }

    #[test] fn parses_revese_slice_test() {
        assert_eq!(parse("[::-1].a").unwrap(),
                   ArrayProjection(Box::new(Slice(None, None, Some(-1))),
                                   Box::new(Identifier("a".to_string()))));
    }

    #[test] fn parses_or_test() {
        assert_eq!(parse("a || b").unwrap(),
                   Or(Box::new(Identifier("a".to_string())),
                      Box::new(Identifier("b".to_string()))));
    }

    #[test] fn parses_pipe_test() {
        assert_eq!(parse("a | b").unwrap(),
                   Subexpr(Box::new(Identifier("a".to_string())),
                      Box::new(Identifier("b".to_string()))));
    }
}
<commit_msg>Adding literal parsing<commit_after>extern crate rustc_serialize;

pub use self::Ast::*;

use lexer::Lexer;
use lexer::Token;
use self::rustc_serialize::json::{Json};

use std::iter::Peekable;

/// Parses a JMESPath expression into an AST
pub fn parse(expr: &str) -> Result<Ast, ParseError> {
    Parser::new(expr).parse()
}

/// Represents the abstract syntax tree of a JMESPath expression.
#[derive(Clone, PartialEq, Debug)]
pub enum Ast {
    Comparison(Comparator, Box<Ast>, Box<Ast>),
    CurrentNode,
    Expref(Box<Ast>),
    Flatten(Box<Ast>),
    Function(char, Vec<(Box<Ast>, Box<Ast>)>),
    Identifier(String),
    Index(i32),
    Literal(Box<Json>),
    MultiList(Vec<Box<Ast>>),
    MultiHash(Vec<Vec<(char, Box<Ast>)>>),
    ArrayProjection(Box<Ast>, Box<Ast>),
    ObjectProjection(Box<Ast>, Box<Ast>),
    Or(Box<Ast>, Box<Ast>),
    Slice(Option<i32>, Option<i32>, Option<i32>),
    Subexpr(Box<Ast>, Box<Ast>),
    WildcardIndex,
    WildcardValues,
}

/// Comparators (i.e., less than, greater than, etc.)
#[derive(Clone, PartialEq, Debug)]
pub enum Comparator { Eq, Lt, Le, Ne, Ge, Gt }

/// Encountered when an invalid JMESPath expression is parsed.
#[derive(Clone, PartialEq, Debug)]
pub struct ParseError {
    /// The error message.
    msg: String,
    /// The line number of the error.
    line: usize,
    /// The column of the error.
    col: usize,
}

/// JMESPath parser. Returns an Ast
pub struct Parser<'a> {
    /// Peekable token stream
    stream: Peekable<Lexer<'a>>,
    /// Expression being parsed
    expr: String,
    /// The current token
    token: Token,
    /// The current character offset in the expression
    pos: usize,
}

impl<'a> Parser<'a> {
    // Constructs a new lexer using the given expression string.
    pub fn new(expr: &'a str) -> Parser<'a> {
        let mut stream = Lexer::new(expr).peekable();
        let tok0 = stream.next().unwrap();
        Parser {
            stream: stream,
            expr: expr.to_string(),
            token: tok0,
            pos: 0,
        }
    }

    /// Parses the expression into result containing an AST or ParseError.
    pub fn parse(&mut self) -> Result<Ast, ParseError> {
        let result = self.expr(0);
        let token = self.stream.next();
        // After parsing the expression, we should have reached the end of the
        // stream.
        if result.is_err()
        || token.is_none()
        || token.unwrap() == Token::Eof {
            result
        } else {
            return self.err("Did not reach token stream EOF".as_ref());
        }
    }

    /// Ensures that the next token in the token stream is one of the pipe
    /// separated token named provided as the edible argument (e.g.,
    /// "Identifier|Eof").
    fn expect(&mut self, edible: &str) -> Result<Ast, ParseError> {
        self.advance();
        // Get the string name of the token.
        let token_name = self.token.token_to_string();
        if edible.contains(&token_name) {
            return Ok(CurrentNode);
        }
        let msg = format!("Expected one of the following tokens: {:?}", edible);
        self.err(&msg)
    }

    /// Advances the cursor position, skipping any whitespace encountered.
    #[inline]
    fn advance(&mut self) {
        self.pos += self.token.size();
        loop {
            let tok = self.stream.next();
            match tok {
                None => break,
                Some(Token::Whitespace) => self.pos += self.token.size(),
                _ => {
                    self.token = tok.unwrap();
                    break
                }
            }
        }
    }

    /// Main parse function of the Pratt parser that parses while RBP < LBP
    pub fn expr(&mut self, rbp: usize) -> Result<Ast, ParseError> {
        // Parse the NUD token.
        let mut left = match self.token.clone() {
            Token::At               => self.nud_at(),
            Token::Identifier(s, _) => self.nud_identifier(s),
            Token::Star             => self.nud_star(),
            Token::Lbracket         => self.nud_lbracket(),
            Token::Flatten          => self.nud_flatten(),
            Token::Literal(v, _)    => self.nud_literal(v),
            // Token::Ampersand        => self.nud_ampersand(),
            // Token::Lbrace           => self.nud_lbrace(),
            // Token::Filter           => self.nud_filter(),
            _ => { return self.err(&"Unexpected NUD token"); }
        };

        // Parse any LED tokens with a higher binding power.
        while rbp < self.token.lbp() {
            left = match self.token {
                Token::Dot      => self.led_dot(left.unwrap()),
                Token::Lbracket => self.led_lbracket(left.unwrap()),
                Token::Flatten  => self.led_flatten(left.unwrap()),
                Token::Or       => self.led_or(left.unwrap()),
                Token::Pipe     => self.led_pipe(left.unwrap()),
                _ => {
                    self.err(&"Unexpected LED token");
                    break;
                }
            };
        }

        left
    }

    /// Returns a formatted ParseError with the given message.
    fn err(&self, msg: &str) -> Result<Ast, ParseError> {
        // Find each new line and create a formatted error message.
        let mut line = 0;
        let mut col = self.pos;
        Err(ParseError {
            msg: format!("Error at {:?} token, {}: {}",
                         self.token, self.pos, msg),
            col: line,
            line: col
        })
    }

    /// Examples: "@"
    fn nud_at(&mut self) -> Result<Ast, ParseError> {
        self.advance();
        Ok(CurrentNode)
    }

    /// Examples: "Foo"
    fn nud_identifier(&mut self, s: String) -> Result<Ast, ParseError> {
        self.advance();
        Ok(Identifier(s))
    }

    /// Examples: "[0]", "[*]", "[a, b]", "[0:1]", etc...
    fn nud_lbracket(&mut self) -> Result<Ast, ParseError> {
        self.advance();
        match self.token {
            Token::Number(_, _) | Token::Colon => self.parse_array_index(),
            Token::Star => {
                if self.stream.peek() != Some(&Token::Rbracket) {
                    return self.parse_multi_list();
                }
                try!(self.expect("Star"));
                self.parse_wildcard_index()
            },
            _ => self.parse_multi_list()
        }
    }

    /// Examples: foo[*], foo[0], foo[:-1], etc.
    fn led_lbracket(&mut self, lhs: Ast) -> Result<Ast, ParseError> {
        try!(self.expect("Number|Colon|Star"));
        match self.token {
            Token::Number(_, _) | Token::Colon => self.parse_array_index(),
            _ => self.parse_wildcard_index()
        }
    }

    fn nud_literal(&mut self, value: Json) -> Result<Ast, ParseError> {
        self.advance();
        Ok(Literal(Box::new(value)))
    }

    /// Examples: "*" (e.g., "* | *" would be a pipe containing two nud stars)
    fn nud_star(&mut self) -> Result<Ast, ParseError> {
        self.advance();
        self.parse_wildcard_values(CurrentNode)
    }

    /// Examples: "[]". Turns it into a LED flatten (i.e., "@[]").
    fn nud_flatten(&mut self) -> Result<Ast, ParseError> {
        self.led_flatten(CurrentNode)
    }

    /// Creates a Projection AST node for a flatten token.
    fn led_flatten(&mut self, lhs: Ast) -> Result<Ast, ParseError> {
        let rhs = try!(self.projection_rhs(Token::Flatten.lbp()));
        Ok(ArrayProjection(
            Box::new(Flatten(Box::new(lhs))),
            Box::new(rhs)
        ))
    }

    fn led_dot(&mut self, left: Ast) -> Result<Ast, ParseError> {
        let rhs = try!(self.parse_dot(Token::Dot.lbp()));
        Ok(Ast::Subexpr(Box::new(left), Box::new(rhs)))
    }

    fn led_or(&mut self, left: Ast) -> Result<Ast, ParseError> {
        self.advance();
        let rhs = try!(self.expr(Token::Or.lbp()));
        Ok(Or(Box::new(left), Box::new(rhs)))
    }

    fn led_pipe(&mut self, left: Ast) -> Result<Ast, ParseError> {
        self.advance();
        let rhs = try!(self.expr(Token::Pipe.lbp()));
        Ok(Subexpr(Box::new(left), Box::new(rhs)))
    }

    /// Parses the right hand side of a dot expression.
    fn parse_dot(&mut self, lbp: usize) -> Result<Ast, ParseError> {
        try!(self.expect("Identifier|Star|Lbrace|Lbracket|Ampersand|Filter"));
        match self.token {
            Token::Lbracket => { self.advance(); self.parse_multi_list() },
            _ => self.expr(lbp)
        }
    }

    /// Parses the right hand side of a projection, using the given LBP to
    /// determine when to stop consuming tokens.
    fn projection_rhs(&mut self, lbp: usize) -> Result<Ast, ParseError> {
        let lbp = self.token.lbp();
        match self.token {
            Token::Dot      => self.parse_dot(lbp),
            Token::Lbracket => self.expr(lbp),
            Token::Filter   => self.expr(lbp),
            _ if lbp < 10   => Ok(CurrentNode),
            _               => self.err("Syntax error found in projection")
        }
    }

    /// Creates a projection for "[*]"
    fn parse_wildcard_index(&mut self) -> Result<Ast, ParseError> {
        try!(self.expect("Rbracket"));
        let lhs = Box::new(CurrentNode);
        let rhs = try!(self.projection_rhs(Token::Star.lbp()));
        Ok(ArrayProjection(lhs, Box::new(rhs)))
    }

    /// Creates a projection for "*"
    fn parse_wildcard_values(&mut self, lhs: Ast) -> Result<Ast, ParseError> {
        let rhs = try!(self.projection_rhs(Token::Star.lbp()));
        Ok(ObjectProjection(Box::new(lhs), Box::new(rhs)))
    }

    /// Parses [0], [::-1], [0:-1], [0:1], etc...
    fn parse_array_index(&mut self) -> Result<Ast, ParseError> {
        let mut parts = [None, None, None];
        let mut pos = 0;
        loop {
            match self.token {
                Token::Colon => {
                    pos += 1;
                    if pos > 2 {
                        return self.err("Too many colons in slice expr");
                    }
                    try!(self.expect("Number|Colon|Rbracket"));
                },
                Token::Number(value, _) => {
                    parts[pos] = Some(value);
                    try!(self.expect("Colon|Rbracket"));
                },
                Token::Rbracket => { self.advance(); break; },
                _ => { return self.err("Unexpected token"); },
            }
        }

        if pos == 0 {
            // No colons were found, so this is a simple index extraction.
            Ok(Index(parts[0].unwrap()))
        } else {
            // Sliced array from start (e.g., [2:])
            let lhs = Slice(parts[0], parts[1], parts[2]);
            let rhs = try!(self.projection_rhs(Token::Star.lbp()));
            Ok(ArrayProjection(Box::new(lhs), Box::new(rhs)))
        }
    }

    /// Parses multi-select lists (e.g., "[foo, bar, baz]")
    fn parse_multi_list(&mut self) -> Result<Ast, ParseError> {
        let mut nodes = vec!();
        loop {
            nodes.push(Box::new(try!(self.expr(0))));
            match self.token {
                // Skip comma tokens
                Token::Comma => {
                    self.advance();
                    // The closing "Rbracket" token cannot follow a comma.
                    if self.token.token_to_string() == "Rbracket" {
                        return self.err("Unexpected token after comma");
                    }
                },
                // Terminal conditon
                Token::Rbracket => break,
                // Got to EOF without closing
                Token::Eof => return self.err("Unclosed multi-select"),
                _ => {}
            }
        }

        Ok(MultiList(nodes))
    }
}

#[cfg(test)]
mod test {
    extern crate rustc_serialize;

    use super::*;
    use self::rustc_serialize::json::{Json};

    #[test] fn indentifier_test() {
        assert_eq!(parse("foo").unwrap(), Identifier("foo".to_string()));
    }

    #[test] fn current_node_test() {
        assert_eq!(parse("@").unwrap(), CurrentNode);
    }

    #[test] fn wildcard_values_test() {
        assert_eq!(parse("*").unwrap(),
                   ObjectProjection(Box::new(CurrentNode),
                                    Box::new(CurrentNode)));
    }

    #[test] fn dot_test() {
        assert!(parse("@.b").unwrap() == Subexpr(Box::new(CurrentNode),
                                                 Box::new(Identifier("b".to_string()))));
    }

    #[test] fn ensures_nud_token_is_valid_test() {

    }

    #[test] fn multi_list_test() {
        let l = MultiList(vec![Box::new(Identifier("a".to_string())),
                               Box::new(Identifier("b".to_string()))]);
        assert_eq!(parse("[a, b]").unwrap(), l);
    }

    #[test] fn multi_list_after_dot_test() {
        let l = MultiList(vec![Box::new(Identifier("a".to_string())),
                               Box::new(Identifier("b".to_string()))]);
        assert_eq!(parse("@.[a, b]").unwrap(), Subexpr(Box::new(CurrentNode), Box::new(l)));
    }

    #[test] fn parses_simple_index_extractions_test() {
        assert_eq!(parse("[0]").unwrap(), Index(0));
    }

    #[test] fn parses_single_element_slice_test() {
        assert_eq!(parse("[-1:]").unwrap(),
                   ArrayProjection(Box::new(Slice(Some(-1), None, None)),
                                   Box::new(CurrentNode)));
    }

    #[test] fn parses_double_element_slice_test() {
        assert_eq!(parse("[1:-1].a").unwrap(),
                   ArrayProjection(Box::new(Slice(Some(1), Some(-1), None)),
                                   Box::new(Identifier("a".to_string()))));
    }

    #[test] fn parses_revese_slice_test() {
        assert_eq!(parse("[::-1].a").unwrap(),
                   ArrayProjection(Box::new(Slice(None, None, Some(-1))),
                                   Box::new(Identifier("a".to_string()))));
    }

    #[test] fn parses_or_test() {
        assert_eq!(parse("a || b").unwrap(),
                   Or(Box::new(Identifier("a".to_string())),
                      Box::new(Identifier("b".to_string()))));
    }

    #[test] fn parses_pipe_test() {
        assert_eq!(parse("a | b").unwrap(),
                   Subexpr(Box::new(Identifier("a".to_string())),
                      Box::new(Identifier("b".to_string()))));
    }

    #[test] fn parses_literal_token_test() {
        assert_eq!(parse("`\"foo\"`").unwrap(),
                   Literal(Box::new(Json::String("foo".to_string()))))
    }
}
<|endoftext|>


<commit_before>/*
* Huffman table generation:
* HuffTable::huff_decode()
* HuffTable::initialize()
* and used data structures are originally grabbed from the IJG software,
* and adapted by Hubert Figuiere.
*
* Copyright (C) 1991, 1992, Thomas G. Lane.
* Part of the Independent JPEG Group's software.
* See the file Copyright for more details.
*
* Copyright (c) 1993 Brian C. Smith, The Regents of the University
* of California
* All rights reserved.
*
* Copyright (c) 1994 Kongji Huang and Brian C. Smith.
* Cornell University
* All rights reserved.
*
* Permission to use, copy, modify, and distribute this software and its
* documentation for any purpose, without fee, and without written agreement is
* hereby granted, provided that the above copyright notice and the following
* two paragraphs appear in all copies of this software.
*
* IN NO EVENT SHALL CORNELL UNIVERSITY BE LIABLE TO ANY PARTY FOR
* DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT
* OF THE USE OF THIS SOFTWARE AND ITS DOCUMENTATION, EVEN IF CORNELL
* UNIVERSITY HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
* CORNELL UNIVERSITY SPECIFICALLY DISCLAIMS ANY WARRANTIES,
* INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY
* AND FITNESS FOR A PARTICULAR PURPOSE.  THE SOFTWARE PROVIDED HEREUNDER IS
* ON AN "AS IS" BASIS, AND CORNELL UNIVERSITY HAS NO OBLIGATION TO
* PROVIDE MAINTENANCE, SUPPORT, UPDATES, ENHANCEMENTS, OR MODIFICATIONS.
*/

use std::fmt;
use crate::decoders::basics::*;

const DECODE_TABLE_BITS: u32 = 13;
const LENGTH_TABLE_BITS: u32 = 8;

pub struct HuffTable {
  // These two fields directly represent the contents of a JPEG DHT marker
  pub bits: [u32;17],
  pub huffval: [u32;256],

  // Represent the weird shifts that are needed for some NEF files
  pub shiftval: [u32;256],

  // The remaining fields are computed from the above to allow more
  // efficient coding and decoding and thus private
  mincode: [u16;17],
  maxcode: [i32;18],
  valptr: [i16;17],
  lengthtable: Vec<Option<(u32,u32,u32)>>,
  decodetable: Vec<Option<(u32,i32)>>,
  pub dng_bug: bool,
  pub initialized: bool,
}

struct MockPump {
  bits: u64,
  nbits: u32,
}

impl MockPump {
  pub fn empty() -> Self {
    MockPump {
      bits: 0,
      nbits: 0,
    }
  }

  pub fn set(&mut self, bits: u32, nbits: u32) {
    self.bits = (bits as u64) << 32;
    self.nbits = nbits + 32;
  }

  pub fn validbits(&self) -> i32 {
    self.nbits as i32 - 32
  }
}

impl BitPump for MockPump {
  fn peek_bits(&mut self, num: u32) -> u32 {
    (self.bits >> (self.nbits-num)) as u32
  }

  fn consume_bits(&mut self, num: u32) {
    self.nbits -= num;
    self.bits &= (1 << self.nbits) - 1;
  }
}

impl HuffTable {
  pub fn empty() -> HuffTable {
    HuffTable {
      bits: [0;17],
      huffval: [0;256],
      shiftval: [0;256],
      mincode: [0;17],
      maxcode: [0;18],
      valptr: [0;17],
      lengthtable: Vec::new(),
      decodetable: Vec::new(),
      dng_bug: false,
      initialized: false,
    }
  }

  pub fn new(bits: [u32;17], huffval: [u32;256], dng_bug: bool) -> Result<HuffTable,String> {
    let mut tbl = HuffTable {
      bits: bits,
      huffval: huffval,
      shiftval: [0;256],
      mincode: [0;17],
      maxcode: [0;18],
      valptr: [0;17],
      lengthtable: Vec::new(),
      decodetable: Vec::new(),
      dng_bug: dng_bug,
      initialized: false,
    };
    tbl.initialize()?;
    Ok(tbl)
  }

  pub fn initialize(&mut self) -> Result<(), String> {
    // Figure C.1: make table of Huffman code length for each symbol
    // Note that this is in code-length order.
    let mut p = 0;
    let mut huffsize: [u8;257] = [0;257];
    for l in 1..17 {
      for _ in 1..((self.bits[l] as usize)+1) {
        huffsize[p] = l as u8;
        p += 1;
        if p > 256 {
          return Err("ljpeg: Code length too long. Corrupt data.".to_string())
        }
      }
    }
    huffsize[p] = 0;

    // Figure C.2: generate the codes themselves
    // Note that this is in code-length order.
    let mut code: u16 = 0;
    let mut si: u32 = huffsize[0] as u32;
    let mut huffcode: [u16;257] = [0;257];
    p = 0;
    while huffsize[p] > 0 {
      while (huffsize[p] as u32) == si {
        huffcode[p] = code;
        p += 1;
        code += 1;
      }
      code <<= 1;
      si += 1;
      if p > 256 {
        return Err("ljpeg: Code length too long. Corrupt data.".to_string())
      }
    }

    //Figure F.15: generate decoding tables
    self.mincode[0] = 0;
    self.maxcode[0] = 0;
    p = 0;
    for l in 1..17 {
      if self.bits[l] > 0 {
        self.valptr[l] = p as i16;
        self.mincode[l] = huffcode[p];
        p += self.bits[l] as usize;
        self.maxcode[l] = huffcode[p - 1] as i32;
      } else {
        self.valptr[l] = 0xff;   // This check must be present to avoid crash on junk
        self.maxcode[l] = -1;
      }
      if p > 256 {
        return Err("ljpeg: Code length too long. Corrupt data.".to_string())
      }
    }

    // We put in this value to ensure HuffDecode terminates
    self.maxcode[17] = 0xFFFFF;

    // Bootstrap the length table with the slow code
    let mut pump = MockPump::empty();
    self.lengthtable = vec![None; 1 << LENGTH_TABLE_BITS];
    let mut i = 0;
    loop {
      pump.set(i, LENGTH_TABLE_BITS);
      let res = self.huff_len_slow(&mut pump);
      let validbits = pump.validbits();
      if validbits >= 0 {
        // We had a valid decode within the lookup bits, save that result to
        // every position where the decode applies.
        for _ in 0..(1 << validbits) {
          self.lengthtable[i as usize] = Some(res);
          i += 1;
        }
      } else {
        i += 1;
      }
      if i >= 1 << LENGTH_TABLE_BITS {
        break;
      }
    }

    // Now bootstrap the full decode table
    let mut pump = MockPump::empty();
    self.decodetable = vec![None; 1 << DECODE_TABLE_BITS];
    let mut i = 0;
    loop {
      pump.set(i, DECODE_TABLE_BITS);
      let res = self.huff_decode_slow(&mut pump);
      let validbits = pump.validbits();
      if validbits >= 0 {
        // We had a valid decode within the lookup bits, save that result to
        // every position where the decode applies.
        for _ in 0..(1 << validbits) {
          self.decodetable[i as usize] = Some(res);
          i += 1;
        }
      } else {
        i += 1;
      }
      if i >= 1 << DECODE_TABLE_BITS {
        break;
      }
    }

    self.initialized = true;
    Ok(())
  }

  pub fn huff_decode(&self, pump: &mut dyn BitPump) -> Result<i32,String> {
    let code = pump.peek_bits(DECODE_TABLE_BITS) as usize;
    if let Some((bits,decode)) = self.decodetable[code] {
      pump.consume_bits(bits);
      Ok(decode)
    } else {
      let decode = self.huff_decode_slow(pump);
      Ok(decode.1)
    }
  }

  pub fn huff_decode_slow(&self, pump: &mut dyn BitPump) -> (u32,i32) {
    let len = self.huff_len(pump);
    (len.0+len.1, self.huff_diff(pump, len))
  }

  pub fn huff_len(&self, pump: &mut dyn BitPump) -> (u32,u32,u32) {
    let code = pump.peek_bits(LENGTH_TABLE_BITS) as usize;
    if let Some((bits,len,shift)) = self.lengthtable[code] {
      pump.consume_bits(bits);
      (bits, len, shift)
    } else {
      self.huff_len_slow(pump)
    }
  }

  pub fn huff_len_slow(&self, pump: &mut dyn BitPump) -> (u32,u32,u32) {
    let mut code = 0 as u32;
    let mut l = 0 as usize;
    loop {
      let temp = pump.get_bits(1);
      code = (code << 1) | temp;
      l += 1;
      if code as i32 <= self.maxcode[l] {
        break;
      }
    }
    let idx = self.valptr[l] as usize + (code as usize - (self.mincode[l] as usize)) as usize;
    (l as u32,self.huffval[idx],self.shiftval[idx])
  }

  pub fn huff_diff(&self, pump: &mut dyn BitPump, input: (u32,u32,u32)) -> i32 {
    let (_, len, shift) = input;

    match len {
      0 => 0,
      16 => {
        if self.dng_bug {
          pump.get_bits(16); // consume can fail because we haven't peeked yet
        }
        -32768
      },
      len => {
        // decode the difference and extend sign bit
        let fulllen: i32 = (len + shift) as i32;
        let shift: i32 = shift as i32;
        let bits = pump.get_bits(len) as i32;
        let mut diff: i32 = ((bits << 1) + 1) << shift >> 1;
        if (diff & (1 << (fulllen - 1))) == 0 {
          diff -= (1 << fulllen) - ((shift == 0) as i32);
        }
        diff
      },
    }
  }
}

impl fmt::Debug for HuffTable {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    if self.initialized {
      write!(f, "HuffTable {{ bits: {:?} huffval: {:?} }}", self.bits, &self.huffval[..])
    } else {
      write!(f, "HuffTable {{ uninitialized }}")
    }
  }
}
<commit_msg>Fix broken huffman decode table<commit_after>/*
* Huffman table generation:
* HuffTable::huff_decode()
* HuffTable::initialize()
* and used data structures are originally grabbed from the IJG software,
* and adapted by Hubert Figuiere.
*
* Copyright (C) 1991, 1992, Thomas G. Lane.
* Part of the Independent JPEG Group's software.
* See the file Copyright for more details.
*
* Copyright (c) 1993 Brian C. Smith, The Regents of the University
* of California
* All rights reserved.
*
* Copyright (c) 1994 Kongji Huang and Brian C. Smith.
* Cornell University
* All rights reserved.
*
* Permission to use, copy, modify, and distribute this software and its
* documentation for any purpose, without fee, and without written agreement is
* hereby granted, provided that the above copyright notice and the following
* two paragraphs appear in all copies of this software.
*
* IN NO EVENT SHALL CORNELL UNIVERSITY BE LIABLE TO ANY PARTY FOR
* DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT
* OF THE USE OF THIS SOFTWARE AND ITS DOCUMENTATION, EVEN IF CORNELL
* UNIVERSITY HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
* CORNELL UNIVERSITY SPECIFICALLY DISCLAIMS ANY WARRANTIES,
* INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY
* AND FITNESS FOR A PARTICULAR PURPOSE.  THE SOFTWARE PROVIDED HEREUNDER IS
* ON AN "AS IS" BASIS, AND CORNELL UNIVERSITY HAS NO OBLIGATION TO
* PROVIDE MAINTENANCE, SUPPORT, UPDATES, ENHANCEMENTS, OR MODIFICATIONS.
*/

use std::fmt;
use crate::decoders::basics::*;

const DECODE_TABLE_BITS: u32 = 13;
const LENGTH_TABLE_BITS: u32 = 8;

pub struct HuffTable {
  // These two fields directly represent the contents of a JPEG DHT marker
  pub bits: [u32;17],
  pub huffval: [u32;256],

  // Represent the weird shifts that are needed for some NEF files
  pub shiftval: [u32;256],

  // The remaining fields are computed from the above to allow more
  // efficient coding and decoding and thus private
  mincode: [u16;17],
  maxcode: [i32;18],
  valptr: [i16;17],
  lengthtable: Vec<Option<(u32,u32,u32)>>,
  decodetable: Vec<Option<(u32,i32)>>,
  pub dng_bug: bool,
  pub initialized: bool,
}

struct MockPump {
  bits: u64,
  nbits: u32,
}

impl MockPump {
  pub fn empty() -> Self {
    MockPump {
      bits: 0,
      nbits: 0,
    }
  }

  pub fn set(&mut self, bits: u32, nbits: u32) {
    self.bits = (bits as u64) << 32;
    self.nbits = nbits + 32;
  }

  pub fn validbits(&self) -> i32 {
    self.nbits as i32 - 32
  }
}

impl BitPump for MockPump {
  fn peek_bits(&mut self, num: u32) -> u32 {
    (self.bits >> (self.nbits-num)) as u32
  }

  fn consume_bits(&mut self, num: u32) {
    self.nbits -= num;
    self.bits &= (1 << self.nbits) - 1;
  }
}

impl HuffTable {
  pub fn empty() -> HuffTable {
    HuffTable {
      bits: [0;17],
      huffval: [0;256],
      shiftval: [0;256],
      mincode: [0;17],
      maxcode: [0;18],
      valptr: [0;17],
      lengthtable: Vec::new(),
      decodetable: Vec::new(),
      dng_bug: false,
      initialized: false,
    }
  }

  pub fn new(bits: [u32;17], huffval: [u32;256], dng_bug: bool) -> Result<HuffTable,String> {
    let mut tbl = HuffTable {
      bits: bits,
      huffval: huffval,
      shiftval: [0;256],
      mincode: [0;17],
      maxcode: [0;18],
      valptr: [0;17],
      lengthtable: Vec::new(),
      decodetable: Vec::new(),
      dng_bug: dng_bug,
      initialized: false,
    };
    tbl.initialize()?;
    Ok(tbl)
  }

  pub fn initialize(&mut self) -> Result<(), String> {
    // Figure C.1: make table of Huffman code length for each symbol
    // Note that this is in code-length order.
    let mut p = 0;
    let mut huffsize: [u8;257] = [0;257];
    for l in 1..17 {
      for _ in 1..((self.bits[l] as usize)+1) {
        huffsize[p] = l as u8;
        p += 1;
        if p > 256 {
          return Err("ljpeg: Code length too long. Corrupt data.".to_string())
        }
      }
    }
    huffsize[p] = 0;

    // Figure C.2: generate the codes themselves
    // Note that this is in code-length order.
    let mut code: u16 = 0;
    let mut si: u32 = huffsize[0] as u32;
    let mut huffcode: [u16;257] = [0;257];
    p = 0;
    while huffsize[p] > 0 {
      while (huffsize[p] as u32) == si {
        huffcode[p] = code;
        p += 1;
        code += 1;
      }
      code <<= 1;
      si += 1;
      if p > 256 {
        return Err("ljpeg: Code length too long. Corrupt data.".to_string())
      }
    }

    //Figure F.15: generate decoding tables
    self.mincode[0] = 0;
    self.maxcode[0] = 0;
    p = 0;
    for l in 1..17 {
      if self.bits[l] > 0 {
        self.valptr[l] = p as i16;
        self.mincode[l] = huffcode[p];
        p += self.bits[l] as usize;
        self.maxcode[l] = huffcode[p - 1] as i32;
      } else {
        self.valptr[l] = 0xff;   // This check must be present to avoid crash on junk
        self.maxcode[l] = -1;
      }
      if p > 256 {
        return Err("ljpeg: Code length too long. Corrupt data.".to_string())
      }
    }

    // We put in this value to ensure HuffDecode terminates
    self.maxcode[17] = 0xFFFFF;

    // Bootstrap the length table with the slow code
    let mut pump = MockPump::empty();
    self.lengthtable = vec![None; 1 << LENGTH_TABLE_BITS];
    let mut i = 0;
    loop {
      pump.set(i, LENGTH_TABLE_BITS);
      let res = self.huff_len_slow(&mut pump);
      let validbits = pump.validbits();
      if validbits >= 0 {
        // We had a valid decode within the lookup bits, save that result to
        // every position where the decode applies.
        for _ in 0..(1 << validbits) {
          self.lengthtable[i as usize] = Some(res);
          i += 1;
        }
      } else {
        i += 1;
      }
      if i >= 1 << LENGTH_TABLE_BITS {
        break;
      }
    }

    // Now bootstrap the full decode table
    let mut pump = MockPump::empty();
    self.decodetable = vec![None; 1 << DECODE_TABLE_BITS];
    let mut i = 0;
    loop {
      pump.set(i, DECODE_TABLE_BITS);
      let res = self.huff_decode_slow(&mut pump);
      let validbits = pump.validbits();
      if validbits >= 0 {
        // We had a valid decode within the lookup bits, save that result to
        // every position where the decode applies.
        for _ in 0..(1 << validbits) {
          self.decodetable[i as usize] = Some(res);
          i += 1;
        }
      } else {
        i += 1;
      }
      if i >= 1 << DECODE_TABLE_BITS {
        break;
      }
    }

    self.initialized = true;
    Ok(())
  }

  pub fn huff_decode(&self, pump: &mut dyn BitPump) -> Result<i32,String> {
    let code = pump.peek_bits(DECODE_TABLE_BITS) as usize;
    if let Some((bits,decode)) = self.decodetable[code] {
      pump.consume_bits(bits);
      Ok(decode)
    } else {
      let decode = self.huff_decode_slow(pump);
      Ok(decode.1)
    }
  }

  pub fn huff_decode_slow(&self, pump: &mut dyn BitPump) -> (u32,i32) {
    let len = self.huff_len(pump);
    (len.0+len.1, self.huff_diff(pump, len))
  }

  pub fn huff_len(&self, pump: &mut dyn BitPump) -> (u32,u32,u32) {
    let code = pump.peek_bits(LENGTH_TABLE_BITS) as usize;
    if let Some((bits,len,shift)) = self.lengthtable[code] {
      pump.consume_bits(bits);
      (bits, len, shift)
    } else {
      self.huff_len_slow(pump)
    }
  }

  pub fn huff_len_slow(&self, pump: &mut dyn BitPump) -> (u32,u32,u32) {
    let mut code = 0 as u32;
    let mut l = 0 as usize;
    loop {
      let temp = pump.get_bits(1);
      code = (code << 1) | temp;
      l += 1;
      if code as i32 <= self.maxcode[l] {
        break;
      }
    }
    if l >= 17 {
      (0, 0, 0)
    } else {
      let idx = self.valptr[l] as usize + (code as usize - (self.mincode[l] as usize)) as usize;
      (l as u32,self.huffval[idx],self.shiftval[idx])
    }
  }

  pub fn huff_diff(&self, pump: &mut dyn BitPump, input: (u32,u32,u32)) -> i32 {
    let (_, len, shift) = input;

    match len {
      0 => 0,
      16 => {
        if self.dng_bug {
          pump.get_bits(16); // consume can fail because we haven't peeked yet
        }
        -32768
      },
      len => {
        // decode the difference and extend sign bit
        let fulllen: i32 = (len + shift) as i32;
        let shift: i32 = shift as i32;
        let bits = pump.get_bits(len) as i32;
        let mut diff: i32 = ((bits << 1) + 1) << shift >> 1;
        if (diff & (1 << (fulllen - 1))) == 0 {
          diff -= (1 << fulllen) - ((shift == 0) as i32);
        }
        diff
      },
    }
  }
}

impl fmt::Debug for HuffTable {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    if self.initialized {
      write!(f, "HuffTable {{ bits: {:?} huffval: {:?} }}", self.bits, &self.huffval[..])
    } else {
      write!(f, "HuffTable {{ uninitialized }}")
    }
  }
}
<|endoftext|>


<commit_before>//! Pragma helpers

use std::ops::Deref;

use crate::error::Error;
use crate::ffi;
use crate::types::{ToSql, ToSqlOutput, ValueRef};
use crate::{Connection, DatabaseName, Result, Row, NO_PARAMS};

pub struct Sql {
    buf: String,
}

impl Sql {
    pub fn new() -> Sql {
        Sql { buf: String::new() }
    }

    pub fn push_pragma(
        &mut self,
        schema_name: Option<DatabaseName<'_>>,
        pragma_name: &str,
    ) -> Result<()> {
        self.push_keyword("PRAGMA")?;
        self.push_space();
        if let Some(schema_name) = schema_name {
            self.push_schema_name(schema_name);
            self.push_dot();
        }
        self.push_keyword(pragma_name)
    }

    pub fn push_keyword(&mut self, keyword: &str) -> Result<()> {
        if !keyword.is_empty() && is_identifier(keyword) {
            self.buf.push_str(keyword);
            Ok(())
        } else {
            Err(Error::SqliteFailure(
                ffi::Error::new(ffi::SQLITE_MISUSE),
                Some(format!("Invalid keyword \"{}\"", keyword)),
            ))
        }
    }

    pub fn push_schema_name(&mut self, schema_name: DatabaseName<'_>) {
        match schema_name {
            DatabaseName::Main => self.buf.push_str("main"),
            DatabaseName::Temp => self.buf.push_str("temp"),
            DatabaseName::Attached(s) => self.push_identifier(s),
        };
    }

    pub fn push_identifier(&mut self, s: &str) {
        if is_identifier(s) {
            self.buf.push_str(s);
        } else {
            self.wrap_and_escape(s, '"');
        }
    }

    pub fn push_value(&mut self, value: &dyn ToSql) -> Result<()> {
        let value = value.to_sql()?;
        let value = match value {
            ToSqlOutput::Borrowed(v) => v,
            ToSqlOutput::Owned(ref v) => ValueRef::from(v),
            #[cfg(feature = "blob")]
            ToSqlOutput::ZeroBlob(_) => {
                return Err(Error::SqliteFailure(
                    ffi::Error::new(ffi::SQLITE_MISUSE),
                    Some(format!("Unsupported value \"{:?}\"", value)),
                ));
            }
            #[cfg(feature = "array")]
            ToSqlOutput::Array(_) => {
                return Err(Error::SqliteFailure(
                    ffi::Error::new(ffi::SQLITE_MISUSE),
                    Some(format!("Unsupported value \"{:?}\"", value)),
                ));
            }
        };
        match value {
            ValueRef::Integer(i) => {
                self.push_int(i);
            }
            ValueRef::Real(r) => {
                self.push_real(r);
            }
            ValueRef::Text(s) => {
                self.push_string_literal(s);
            }
            _ => {
                return Err(Error::SqliteFailure(
                    ffi::Error::new(ffi::SQLITE_MISUSE),
                    Some(format!("Unsupported value \"{:?}\"", value)),
                ));
            }
        };
        Ok(())
    }

    pub fn push_string_literal(&mut self, s: &str) {
        self.wrap_and_escape(s, '\'');
    }

    pub fn push_int(&mut self, i: i64) {
        self.buf.push_str(&i.to_string());
    }

    pub fn push_real(&mut self, f: f64) {
        self.buf.push_str(&f.to_string());
    }

    pub fn push_space(&mut self) {
        self.buf.push(' ');
    }

    pub fn push_dot(&mut self) {
        self.buf.push('.');
    }

    pub fn push_equal_sign(&mut self) {
        self.buf.push('=');
    }

    pub fn open_brace(&mut self) {
        self.buf.push('(');
    }

    pub fn close_brace(&mut self) {
        self.buf.push(')');
    }

    pub fn as_str(&self) -> &str {
        &self.buf
    }

    fn wrap_and_escape(&mut self, s: &str, quote: char) {
        self.buf.push(quote);
        let chars = s.chars();
        for ch in chars {
            // escape `quote` by doubling it
            if ch == quote {
                self.buf.push(ch);
            }
            self.buf.push(ch)
        }
        self.buf.push(quote);
    }
}

impl Deref for Sql {
    type Target = str;

    fn deref(&self) -> &str {
        self.as_str()
    }
}

impl Connection {
    /// Query the current value of `pragma_name`.
    ///
    /// Some pragmas will return multiple rows/values which cannot be retrieved
    /// with this method.
    ///
    /// Prefer [PRAGMA function](https://sqlite.org/pragma.html#pragfunc) introduced in SQLite 3.20:
    /// `SELECT user_version FROM pragma_user_version;`
    pub fn pragma_query_value<T, F>(
        &self,
        schema_name: Option<DatabaseName<'_>>,
        pragma_name: &str,
        f: F,
    ) -> Result<T>
    where
        F: FnOnce(Row<'_, '_>) -> Result<T>,
    {
        let mut query = Sql::new();
        query.push_pragma(schema_name, pragma_name)?;
        let mut stmt = self.prepare(&query)?;
        let mut rows = stmt.query(NO_PARAMS)?;
        if let Some(result_row) = rows.next() {
            let row = result_row?;
            f(row)
        } else {
            Err(Error::QueryReturnedNoRows)
        }
    }

    /// Query the current rows/values of `pragma_name`.
    ///
    /// Prefer [PRAGMA function](https://sqlite.org/pragma.html#pragfunc) introduced in SQLite 3.20:
    /// `SELECT * FROM pragma_collation_list;`
    pub fn pragma_query<F>(
        &self,
        schema_name: Option<DatabaseName<'_>>,
        pragma_name: &str,
        mut f: F,
    ) -> Result<()>
    where
        F: FnMut(&Row<'_, '_>) -> Result<()>,
    {
        let mut query = Sql::new();
        query.push_pragma(schema_name, pragma_name)?;
        let mut stmt = self.prepare(&query)?;
        let mut rows = stmt.query(NO_PARAMS)?;
        while let Some(result_row) = rows.next() {
            let row = result_row?;
            f(&row)?;
        }
        Ok(())
    }

    /// Query the current value(s) of `pragma_name` associated to
    /// `pragma_value`.
    ///
    /// This method can be used with query-only pragmas which need an argument
    /// (e.g. `table_info('one_tbl')`) or pragmas which returns value(s)
    /// (e.g. `integrity_check`).
    ///
    /// Prefer [PRAGMA function](https://sqlite.org/pragma.html#pragfunc) introduced in SQLite 3.20:
    /// `SELECT * FROM pragma_table_info(?);`
    pub fn pragma<F>(
        &self,
        schema_name: Option<DatabaseName<'_>>,
        pragma_name: &str,
        pragma_value: &dyn ToSql,
        mut f: F,
    ) -> Result<()>
    where
        F: FnMut(&Row<'_, '_>) -> Result<()>,
    {
        let mut sql = Sql::new();
        sql.push_pragma(schema_name, pragma_name)?;
        // The argument is may be either in parentheses
        // or it may be separated from the pragma name by an equal sign.
        // The two syntaxes yield identical results.
        sql.open_brace();
        sql.push_value(pragma_value)?;
        sql.close_brace();
        let mut stmt = self.prepare(&sql)?;
        let mut rows = stmt.query(NO_PARAMS)?;
        while let Some(result_row) = rows.next() {
            let row = result_row?;
            f(&row)?;
        }
        Ok(())
    }

    /// Set a new value to `pragma_name`.
    ///
    /// Some pragmas will return the updated value which cannot be retrieved
    /// with this method.
    pub fn pragma_update(
        &self,
        schema_name: Option<DatabaseName<'_>>,
        pragma_name: &str,
        pragma_value: &dyn ToSql,
    ) -> Result<()> {
        let mut sql = Sql::new();
        sql.push_pragma(schema_name, pragma_name)?;
        // The argument is may be either in parentheses
        // or it may be separated from the pragma name by an equal sign.
        // The two syntaxes yield identical results.
        sql.push_equal_sign();
        sql.push_value(pragma_value)?;
        self.execute_batch(&sql)
    }

    /// Set a new value to `pragma_name` and return the updated value.
    ///
    /// Only few pragmas automatically return the updated value.
    pub fn pragma_update_and_check<F, T>(
        &self,
        schema_name: Option<DatabaseName<'_>>,
        pragma_name: &str,
        pragma_value: &dyn ToSql,
        f: F,
    ) -> Result<T>
    where
        F: FnOnce(Row<'_, '_>) -> Result<T>,
    {
        let mut sql = Sql::new();
        sql.push_pragma(schema_name, pragma_name)?;
        // The argument is may be either in parentheses
        // or it may be separated from the pragma name by an equal sign.
        // The two syntaxes yield identical results.
        sql.push_equal_sign();
        sql.push_value(pragma_value)?;
        let mut stmt = self.prepare(&sql)?;
        let mut rows = stmt.query(NO_PARAMS)?;
        if let Some(result_row) = rows.next() {
            let row = result_row?;
            f(row)
        } else {
            Err(Error::QueryReturnedNoRows)
        }
    }
}

fn is_identifier(s: &str) -> bool {
    let chars = s.char_indices();
    for (i, ch) in chars {
        if i == 0 {
            if !is_identifier_start(ch) {
                return false;
            }
        } else if !is_identifier_continue(ch) {
            return false;
        }
    }
    true
}

fn is_identifier_start(c: char) -> bool {
    (c >= 'A' && c <= 'Z') || c == '_' || (c >= 'a' && c <= 'z') || c > '\x7F'
}

fn is_identifier_continue(c: char) -> bool {
    c == '$'
        || (c >= '0' && c <= '9')
        || (c >= 'A' && c <= 'Z')
        || c == '_'
        || (c >= 'a' && c <= 'z')
        || c > '\x7F'
}

#[cfg(test)]
mod test {
    use super::Sql;
    use crate::pragma;
    use crate::{Connection, DatabaseName, NO_PARAMS};

    #[test]
    fn pragma_query_value() {
        let db = Connection::open_in_memory().unwrap();
        let user_version: i32 = db
            .pragma_query_value(None, "user_version", |row| row.get_checked(0))
            .unwrap();
        assert_eq!(0, user_version);
    }

    #[test]
    #[cfg(feature = "bundled")]
    fn pragma_func_query_value() {
        let db = Connection::open_in_memory().unwrap();
        let user_version: i32 = db
            .query_row(
                "SELECT user_version FROM pragma_user_version",
                NO_PARAMS,
                |row| row.get(0),
            )
            .unwrap();
        assert_eq!(0, user_version);
    }

    #[test]
    fn pragma_query_no_schema() {
        let db = Connection::open_in_memory().unwrap();
        let mut user_version = -1;
        db.pragma_query(None, "user_version", |row| {
            user_version = row.get_checked(0)?;
            Ok(())
        })
        .unwrap();
        assert_eq!(0, user_version);
    }

    #[test]
    fn pragma_query_with_schema() {
        let db = Connection::open_in_memory().unwrap();
        let mut user_version = -1;
        db.pragma_query(Some(DatabaseName::Main), "user_version", |row| {
            user_version = row.get_checked(0)?;
            Ok(())
        })
        .unwrap();
        assert_eq!(0, user_version);
    }

    #[test]
    fn pragma() {
        let db = Connection::open_in_memory().unwrap();
        let mut columns = Vec::new();
        db.pragma(None, "table_info", &"sqlite_master", |row| {
            let column: String = row.get_checked(1)?;
            columns.push(column);
            Ok(())
        })
        .unwrap();
        assert_eq!(5, columns.len());
    }

    #[test]
    #[cfg(feature = "bundled")]
    fn pragma_func() {
        let db = Connection::open_in_memory().unwrap();
        let mut table_info = db.prepare("SELECT * FROM pragma_table_info(?)").unwrap();
        let mut columns = Vec::new();
        let mut rows = table_info.query(&["sqlite_master"]).unwrap();

        while let Some(row) = rows.next() {
            let row = row.unwrap();
            let column: String = row.get(1);
            columns.push(column);
        }
        assert_eq!(5, columns.len());
    }

    #[test]
    fn pragma_update() {
        let db = Connection::open_in_memory().unwrap();
        db.pragma_update(None, "user_version", &1).unwrap();
    }

    #[test]
    fn pragma_update_and_check() {
        let db = Connection::open_in_memory().unwrap();
        let journal_mode: String = db
            .pragma_update_and_check(None, "journal_mode", &"OFF", |row| row.get_checked(0))
            .unwrap();
        assert_eq!("off", &journal_mode);
    }

    #[test]
    fn is_identifier() {
        assert!(pragma::is_identifier("full"));
        assert!(pragma::is_identifier("r2d2"));
        assert!(!pragma::is_identifier("sp ce"));
        assert!(!pragma::is_identifier("semi;colon"));
    }

    #[test]
    fn double_quote() {
        let mut sql = Sql::new();
        sql.push_schema_name(DatabaseName::Attached(r#"schema";--"#));
        assert_eq!(r#""schema"";--""#, sql.as_str());
    }

    #[test]
    fn wrap_and_escape() {
        let mut sql = Sql::new();
        sql.push_string_literal("value'; --");
        assert_eq!("'value''; --'", sql.as_str());
    }
}
<commit_msg>Fix warning<commit_after>//! Pragma helpers

use std::ops::Deref;

use crate::error::Error;
use crate::ffi;
use crate::types::{ToSql, ToSqlOutput, ValueRef};
use crate::{Connection, DatabaseName, Result, Row, NO_PARAMS};

pub struct Sql {
    buf: String,
}

impl Sql {
    pub fn new() -> Sql {
        Sql { buf: String::new() }
    }

    pub fn push_pragma(
        &mut self,
        schema_name: Option<DatabaseName<'_>>,
        pragma_name: &str,
    ) -> Result<()> {
        self.push_keyword("PRAGMA")?;
        self.push_space();
        if let Some(schema_name) = schema_name {
            self.push_schema_name(schema_name);
            self.push_dot();
        }
        self.push_keyword(pragma_name)
    }

    pub fn push_keyword(&mut self, keyword: &str) -> Result<()> {
        if !keyword.is_empty() && is_identifier(keyword) {
            self.buf.push_str(keyword);
            Ok(())
        } else {
            Err(Error::SqliteFailure(
                ffi::Error::new(ffi::SQLITE_MISUSE),
                Some(format!("Invalid keyword \"{}\"", keyword)),
            ))
        }
    }

    pub fn push_schema_name(&mut self, schema_name: DatabaseName<'_>) {
        match schema_name {
            DatabaseName::Main => self.buf.push_str("main"),
            DatabaseName::Temp => self.buf.push_str("temp"),
            DatabaseName::Attached(s) => self.push_identifier(s),
        };
    }

    pub fn push_identifier(&mut self, s: &str) {
        if is_identifier(s) {
            self.buf.push_str(s);
        } else {
            self.wrap_and_escape(s, '"');
        }
    }

    pub fn push_value(&mut self, value: &dyn ToSql) -> Result<()> {
        let value = value.to_sql()?;
        let value = match value {
            ToSqlOutput::Borrowed(v) => v,
            ToSqlOutput::Owned(ref v) => ValueRef::from(v),
            #[cfg(feature = "blob")]
            ToSqlOutput::ZeroBlob(_) => {
                return Err(Error::SqliteFailure(
                    ffi::Error::new(ffi::SQLITE_MISUSE),
                    Some(format!("Unsupported value \"{:?}\"", value)),
                ));
            }
            #[cfg(feature = "array")]
            ToSqlOutput::Array(_) => {
                return Err(Error::SqliteFailure(
                    ffi::Error::new(ffi::SQLITE_MISUSE),
                    Some(format!("Unsupported value \"{:?}\"", value)),
                ));
            }
        };
        match value {
            ValueRef::Integer(i) => {
                self.push_int(i);
            }
            ValueRef::Real(r) => {
                self.push_real(r);
            }
            ValueRef::Text(s) => {
                self.push_string_literal(s);
            }
            _ => {
                return Err(Error::SqliteFailure(
                    ffi::Error::new(ffi::SQLITE_MISUSE),
                    Some(format!("Unsupported value \"{:?}\"", value)),
                ));
            }
        };
        Ok(())
    }

    pub fn push_string_literal(&mut self, s: &str) {
        self.wrap_and_escape(s, '\'');
    }

    pub fn push_int(&mut self, i: i64) {
        self.buf.push_str(&i.to_string());
    }

    pub fn push_real(&mut self, f: f64) {
        self.buf.push_str(&f.to_string());
    }

    pub fn push_space(&mut self) {
        self.buf.push(' ');
    }

    pub fn push_dot(&mut self) {
        self.buf.push('.');
    }

    pub fn push_equal_sign(&mut self) {
        self.buf.push('=');
    }

    pub fn open_brace(&mut self) {
        self.buf.push('(');
    }

    pub fn close_brace(&mut self) {
        self.buf.push(')');
    }

    pub fn as_str(&self) -> &str {
        &self.buf
    }

    fn wrap_and_escape(&mut self, s: &str, quote: char) {
        self.buf.push(quote);
        let chars = s.chars();
        for ch in chars {
            // escape `quote` by doubling it
            if ch == quote {
                self.buf.push(ch);
            }
            self.buf.push(ch)
        }
        self.buf.push(quote);
    }
}

impl Deref for Sql {
    type Target = str;

    fn deref(&self) -> &str {
        self.as_str()
    }
}

impl Connection {
    /// Query the current value of `pragma_name`.
    ///
    /// Some pragmas will return multiple rows/values which cannot be retrieved
    /// with this method.
    ///
    /// Prefer [PRAGMA function](https://sqlite.org/pragma.html#pragfunc) introduced in SQLite 3.20:
    /// `SELECT user_version FROM pragma_user_version;`
    pub fn pragma_query_value<T, F>(
        &self,
        schema_name: Option<DatabaseName<'_>>,
        pragma_name: &str,
        f: F,
    ) -> Result<T>
    where
        F: FnOnce(Row<'_, '_>) -> Result<T>,
    {
        let mut query = Sql::new();
        query.push_pragma(schema_name, pragma_name)?;
        let mut stmt = self.prepare(&query)?;
        let mut rows = stmt.query(NO_PARAMS)?;
        if let Some(result_row) = rows.next() {
            let row = result_row?;
            f(row)
        } else {
            Err(Error::QueryReturnedNoRows)
        }
    }

    /// Query the current rows/values of `pragma_name`.
    ///
    /// Prefer [PRAGMA function](https://sqlite.org/pragma.html#pragfunc) introduced in SQLite 3.20:
    /// `SELECT * FROM pragma_collation_list;`
    pub fn pragma_query<F>(
        &self,
        schema_name: Option<DatabaseName<'_>>,
        pragma_name: &str,
        mut f: F,
    ) -> Result<()>
    where
        F: FnMut(&Row<'_, '_>) -> Result<()>,
    {
        let mut query = Sql::new();
        query.push_pragma(schema_name, pragma_name)?;
        let mut stmt = self.prepare(&query)?;
        let mut rows = stmt.query(NO_PARAMS)?;
        while let Some(result_row) = rows.next() {
            let row = result_row?;
            f(&row)?;
        }
        Ok(())
    }

    /// Query the current value(s) of `pragma_name` associated to
    /// `pragma_value`.
    ///
    /// This method can be used with query-only pragmas which need an argument
    /// (e.g. `table_info('one_tbl')`) or pragmas which returns value(s)
    /// (e.g. `integrity_check`).
    ///
    /// Prefer [PRAGMA function](https://sqlite.org/pragma.html#pragfunc) introduced in SQLite 3.20:
    /// `SELECT * FROM pragma_table_info(?);`
    pub fn pragma<F>(
        &self,
        schema_name: Option<DatabaseName<'_>>,
        pragma_name: &str,
        pragma_value: &dyn ToSql,
        mut f: F,
    ) -> Result<()>
    where
        F: FnMut(&Row<'_, '_>) -> Result<()>,
    {
        let mut sql = Sql::new();
        sql.push_pragma(schema_name, pragma_name)?;
        // The argument is may be either in parentheses
        // or it may be separated from the pragma name by an equal sign.
        // The two syntaxes yield identical results.
        sql.open_brace();
        sql.push_value(pragma_value)?;
        sql.close_brace();
        let mut stmt = self.prepare(&sql)?;
        let mut rows = stmt.query(NO_PARAMS)?;
        while let Some(result_row) = rows.next() {
            let row = result_row?;
            f(&row)?;
        }
        Ok(())
    }

    /// Set a new value to `pragma_name`.
    ///
    /// Some pragmas will return the updated value which cannot be retrieved
    /// with this method.
    pub fn pragma_update(
        &self,
        schema_name: Option<DatabaseName<'_>>,
        pragma_name: &str,
        pragma_value: &dyn ToSql,
    ) -> Result<()> {
        let mut sql = Sql::new();
        sql.push_pragma(schema_name, pragma_name)?;
        // The argument is may be either in parentheses
        // or it may be separated from the pragma name by an equal sign.
        // The two syntaxes yield identical results.
        sql.push_equal_sign();
        sql.push_value(pragma_value)?;
        self.execute_batch(&sql)
    }

    /// Set a new value to `pragma_name` and return the updated value.
    ///
    /// Only few pragmas automatically return the updated value.
    pub fn pragma_update_and_check<F, T>(
        &self,
        schema_name: Option<DatabaseName<'_>>,
        pragma_name: &str,
        pragma_value: &dyn ToSql,
        f: F,
    ) -> Result<T>
    where
        F: FnOnce(Row<'_, '_>) -> Result<T>,
    {
        let mut sql = Sql::new();
        sql.push_pragma(schema_name, pragma_name)?;
        // The argument is may be either in parentheses
        // or it may be separated from the pragma name by an equal sign.
        // The two syntaxes yield identical results.
        sql.push_equal_sign();
        sql.push_value(pragma_value)?;
        let mut stmt = self.prepare(&sql)?;
        let mut rows = stmt.query(NO_PARAMS)?;
        if let Some(result_row) = rows.next() {
            let row = result_row?;
            f(row)
        } else {
            Err(Error::QueryReturnedNoRows)
        }
    }
}

fn is_identifier(s: &str) -> bool {
    let chars = s.char_indices();
    for (i, ch) in chars {
        if i == 0 {
            if !is_identifier_start(ch) {
                return false;
            }
        } else if !is_identifier_continue(ch) {
            return false;
        }
    }
    true
}

fn is_identifier_start(c: char) -> bool {
    (c >= 'A' && c <= 'Z') || c == '_' || (c >= 'a' && c <= 'z') || c > '\x7F'
}

fn is_identifier_continue(c: char) -> bool {
    c == '$'
        || (c >= '0' && c <= '9')
        || (c >= 'A' && c <= 'Z')
        || c == '_'
        || (c >= 'a' && c <= 'z')
        || c > '\x7F'
}

#[cfg(test)]
mod test {
    use super::Sql;
    use crate::pragma;
    use crate::{Connection, DatabaseName};

    #[test]
    fn pragma_query_value() {
        let db = Connection::open_in_memory().unwrap();
        let user_version: i32 = db
            .pragma_query_value(None, "user_version", |row| row.get_checked(0))
            .unwrap();
        assert_eq!(0, user_version);
    }

    #[test]
    #[cfg(feature = "bundled")]
    fn pragma_func_query_value() {
        use crate::NO_PARAMS;

        let db = Connection::open_in_memory().unwrap();
        let user_version: i32 = db
            .query_row(
                "SELECT user_version FROM pragma_user_version",
                NO_PARAMS,
                |row| row.get(0),
            )
            .unwrap();
        assert_eq!(0, user_version);
    }

    #[test]
    fn pragma_query_no_schema() {
        let db = Connection::open_in_memory().unwrap();
        let mut user_version = -1;
        db.pragma_query(None, "user_version", |row| {
            user_version = row.get_checked(0)?;
            Ok(())
        })
        .unwrap();
        assert_eq!(0, user_version);
    }

    #[test]
    fn pragma_query_with_schema() {
        let db = Connection::open_in_memory().unwrap();
        let mut user_version = -1;
        db.pragma_query(Some(DatabaseName::Main), "user_version", |row| {
            user_version = row.get_checked(0)?;
            Ok(())
        })
        .unwrap();
        assert_eq!(0, user_version);
    }

    #[test]
    fn pragma() {
        let db = Connection::open_in_memory().unwrap();
        let mut columns = Vec::new();
        db.pragma(None, "table_info", &"sqlite_master", |row| {
            let column: String = row.get_checked(1)?;
            columns.push(column);
            Ok(())
        })
        .unwrap();
        assert_eq!(5, columns.len());
    }

    #[test]
    #[cfg(feature = "bundled")]
    fn pragma_func() {
        let db = Connection::open_in_memory().unwrap();
        let mut table_info = db.prepare("SELECT * FROM pragma_table_info(?)").unwrap();
        let mut columns = Vec::new();
        let mut rows = table_info.query(&["sqlite_master"]).unwrap();

        while let Some(row) = rows.next() {
            let row = row.unwrap();
            let column: String = row.get(1);
            columns.push(column);
        }
        assert_eq!(5, columns.len());
    }

    #[test]
    fn pragma_update() {
        let db = Connection::open_in_memory().unwrap();
        db.pragma_update(None, "user_version", &1).unwrap();
    }

    #[test]
    fn pragma_update_and_check() {
        let db = Connection::open_in_memory().unwrap();
        let journal_mode: String = db
            .pragma_update_and_check(None, "journal_mode", &"OFF", |row| row.get_checked(0))
            .unwrap();
        assert_eq!("off", &journal_mode);
    }

    #[test]
    fn is_identifier() {
        assert!(pragma::is_identifier("full"));
        assert!(pragma::is_identifier("r2d2"));
        assert!(!pragma::is_identifier("sp ce"));
        assert!(!pragma::is_identifier("semi;colon"));
    }

    #[test]
    fn double_quote() {
        let mut sql = Sql::new();
        sql.push_schema_name(DatabaseName::Attached(r#"schema";--"#));
        assert_eq!(r#""schema"";--""#, sql.as_str());
    }

    #[test]
    fn wrap_and_escape() {
        let mut sql = Sql::new();
        sql.push_string_literal("value'; --");
        assert_eq!("'value''; --'", sql.as_str());
    }
}
<|endoftext|>


<commit_before>use std::env;
use std::fs::{File, read_dir};
use std::io::Write;
use std::path::PathBuf;
use std::result::Result;
use serde::{Deserialize, Serialize};
use serde_json::ser::Serializer;
use serde_protobuf::descriptor::{Descriptors, MessageDescriptor};
use serde_protobuf::de::Deserializer;
use serde_protobuf::error::{Error, ErrorKind};
use serde_value::Value;
use protobuf::{CodedInputStream, parse_from_reader};
use protobuf::error::ProtobufError;

pub enum PqrsError {
    EofError(String),
    SerdeError(String),
    ProtobufError(String),
}

pub fn process_single(data: &[u8], msg_type: &str, out: &mut Write) -> Result<(), PqrsError> {
    match msg_type {
        "guess" => return guess_message(data, out),
        _ => return named_message(data, msg_type, out),
    }
}

fn named_message(data: &[u8], msg_type: &str, out: &mut Write) -> Result<(), PqrsError> {
    let mut loc_msg_type = String::new();
    let ch = msg_type.chars().nth(0).unwrap();
    if ch != '.' {
        loc_msg_type.push('.');
    }
    loc_msg_type.push_str(msg_type);

    let (descriptors, _) = get_descriptors(false);

    let stream = CodedInputStream::from_bytes(data);
    let mut deserializer = Deserializer::for_named_message(&descriptors, &loc_msg_type, stream).unwrap();
    let mut serializer = Serializer::new(out);
    match deser(&mut deserializer) {
        Ok(value) => value.serialize(&mut serializer).unwrap(),
        Err(e) => return Err(e),
    }
    Ok(())
}

fn guess_message(data: &[u8], out: &mut Write) -> Result<(), PqrsError> {
    let (descriptors, message_descriptors) = get_descriptors(true);

    let mut serializer = Serializer::new(out);
    for md in message_descriptors {
        let stream = CodedInputStream::from_bytes(data);
        let mut deserializer = Deserializer::new(&descriptors, &md, stream);
        match deser(&mut deserializer) {
            Ok(value) => value.serialize(&mut serializer).unwrap(),
            Err(e) => return Err(e),
        }
    }
    Ok(())
}

fn deser(deserializer: &mut Deserializer) -> Result<Value, PqrsError> {
    match Value::deserialize(deserializer) {
        Ok(x) => return Ok(x),
        Err(Error(ErrorKind::Protobuf(ProtobufError::WireError(msg)), _)) => {
            if msg == "unexpected EOF" {
                return Err(PqrsError::EofError(msg));
            }
            return Err(PqrsError::ProtobufError(msg));
        },
        Err(e) => return Err(PqrsError::SerdeError(String::from(e.description()))),
    };
}

fn get_descriptors(with_message_descriptors: bool) -> (Descriptors, Vec<MessageDescriptor>) {
    let mut descriptors = Descriptors::new();
    let mut message_descriptors = Vec::new();

    for fdset_path in discover_fdsets() {
        let mut fdset_file = File::open(fdset_path.as_path()).unwrap();
        let fdset_proto = parse_from_reader(&mut fdset_file).unwrap();
        descriptors.add_file_set_proto(&fdset_proto);
        if with_message_descriptors {
            for file_proto in fdset_proto.get_file().iter() {
                for message_proto in file_proto.get_message_type().iter() {
                    message_descriptors.push(MessageDescriptor::from_proto(&fdset_path.to_string_lossy().into_owned().as_str(), message_proto));
                }
            }
        }
    }
    descriptors.resolve_refs();
    (descriptors, message_descriptors)
}

fn discover_fdsets() -> Vec<PathBuf> {
    let mut fdset_files = Vec::new();

    let mut home = env::home_dir().expect("Could not find $HOME");
    home.push(".pq");
    let paths = read_dir(home.as_path()).unwrap();

    for p in paths {
        let path = match p {
            Ok(p) => p.path(),
            Err(_) => continue,
        };
        match path.extension() {
            Some(x) => {
                if x != "fdset" {
                    continue;
                }
            },
            None => continue,
        }
        fdset_files.push(path);
    }
    fdset_files
}
<commit_msg>Shorten discover_fdsets function<commit_after>use std::env;
use std::fs::{File, read_dir};
use std::io::Write;
use std::path::PathBuf;
use std::result::Result;
use serde::{Deserialize, Serialize};
use serde_json::ser::Serializer;
use serde_protobuf::descriptor::{Descriptors, MessageDescriptor};
use serde_protobuf::de::Deserializer;
use serde_protobuf::error::{Error, ErrorKind};
use serde_value::Value;
use protobuf::{CodedInputStream, parse_from_reader};
use protobuf::error::ProtobufError;

pub enum PqrsError {
    EofError(String),
    SerdeError(String),
    ProtobufError(String),
}

pub fn process_single(data: &[u8], msg_type: &str, out: &mut Write) -> Result<(), PqrsError> {
    match msg_type {
        "guess" => return guess_message(data, out),
        _ => return named_message(data, msg_type, out),
    }
}

fn named_message(data: &[u8], msg_type: &str, out: &mut Write) -> Result<(), PqrsError> {
    let mut loc_msg_type = String::new();
    let ch = msg_type.chars().nth(0).unwrap();
    if ch != '.' {
        loc_msg_type.push('.');
    }
    loc_msg_type.push_str(msg_type);

    let (descriptors, _) = get_descriptors(false);

    let stream = CodedInputStream::from_bytes(data);
    let mut deserializer = Deserializer::for_named_message(&descriptors, &loc_msg_type, stream).unwrap();
    let mut serializer = Serializer::new(out);
    match deser(&mut deserializer) {
        Ok(value) => value.serialize(&mut serializer).unwrap(),
        Err(e) => return Err(e),
    }
    Ok(())
}

fn guess_message(data: &[u8], out: &mut Write) -> Result<(), PqrsError> {
    let (descriptors, message_descriptors) = get_descriptors(true);

    let mut serializer = Serializer::new(out);
    for md in message_descriptors {
        let stream = CodedInputStream::from_bytes(data);
        let mut deserializer = Deserializer::new(&descriptors, &md, stream);
        match deser(&mut deserializer) {
            Ok(value) => {
                println!("{:#?}", value);
                value.serialize(&mut serializer).unwrap();
            }
            Err(e) => return Err(e),
        }
    }
    Ok(())
}

fn deser(deserializer: &mut Deserializer) -> Result<Value, PqrsError> {
    match Value::deserialize(deserializer) {
        Ok(x) => return Ok(x),
        Err(Error(ErrorKind::Protobuf(ProtobufError::WireError(msg)), _)) => {
            if msg == "unexpected EOF" {
                return Err(PqrsError::EofError(msg));
            }
            return Err(PqrsError::ProtobufError(msg));
        },
        Err(e) => return Err(PqrsError::SerdeError(String::from(e.description()))),
    };
}

fn get_descriptors(with_message_descriptors: bool) -> (Descriptors, Vec<MessageDescriptor>) {
    let mut descriptors = Descriptors::new();
    let mut message_descriptors = Vec::new();

    for fdset_path in discover_fdsets() {
        let mut fdset_file = File::open(fdset_path.as_path()).unwrap();
        let fdset_proto = parse_from_reader(&mut fdset_file).unwrap();
        descriptors.add_file_set_proto(&fdset_proto);
        if with_message_descriptors {
            for file_proto in fdset_proto.get_file().iter() {
                for message_proto in file_proto.get_message_type().iter() {
                    message_descriptors.push(MessageDescriptor::from_proto(&fdset_path.to_string_lossy().into_owned().as_str(), message_proto));
                }
            }
        }
    }
    descriptors.resolve_refs();
    (descriptors, message_descriptors)
}

fn discover_fdsets() -> Vec<PathBuf> {
    let mut home = env::home_dir().expect("Could not find $HOME");
    home.push(".pq");

    read_dir(home.as_path()).unwrap()
        .map(|x| {
            x.unwrap().path()
        })
        .collect::<Vec<_>>()
}
<|endoftext|>


<commit_before>use std::io::{Buffer, Reader, IoError, IoResult, OtherIoError};
use serialize::Decoder;

pub struct DecoderReader<'a, R: 'a> {
    reader: &'a mut R
}

impl<'a, R: Reader+Buffer> DecoderReader<'a, R> {
    pub fn new(r: &'a mut R) -> DecoderReader<'a, R> {
        DecoderReader {
            reader: r
        }
    }
}

impl<'a, R: Reader+Buffer> Decoder<IoError> for DecoderReader<'a, R> {
    fn read_nil(&mut self) -> IoResult<()> {
        Ok(())
    }
    fn read_uint(&mut self) -> IoResult<uint> {
        self.read_u64().map(|x| x as uint)
    }
    fn read_u64(&mut self) -> IoResult<u64> {
        self.reader.read_be_u64()
    }
    fn read_u32(&mut self) -> IoResult<u32> {
        self.reader.read_be_u32()
    }
    fn read_u16(&mut self) -> IoResult<u16> {
        self.reader.read_be_u16()
    }
    fn read_u8(&mut self) -> IoResult<u8> {
        self.reader.read_u8()
    }
    fn read_int(&mut self) -> IoResult<int> {
        self.read_i64().map(|x| x as int)
    }
    fn read_i64(&mut self) -> IoResult<i64> {
        self.reader.read_be_i64()
    }
    fn read_i32(&mut self) -> IoResult<i32> {
        self.reader.read_be_i32()
    }
    fn read_i16(&mut self) -> IoResult<i16> {
        self.reader.read_be_i16()
    }
    fn read_i8(&mut self) -> IoResult<i8> {
        self.reader.read_i8()
    }
    fn read_bool(&mut self) -> IoResult<bool> {
        match try!(self.reader.read_i8()) {
            1 => Ok(true),
            _ => Ok(false)
        }
    }
    fn read_f64(&mut self) -> IoResult<f64> {
        self.reader.read_be_f64()
    }
    fn read_f32(&mut self) -> IoResult<f32> {
        self.reader.read_be_f32()
    }
    fn read_char(&mut self) -> IoResult<char> {
        self.reader.read_char()
    }
    fn read_str(&mut self) -> IoResult<String> {
        let len = try!(self.read_uint());
        let mut vector = Vec::with_capacity(len as uint);
        for _ in range(0, len) {
            vector.push(try!(self.reader.read_u8()));
        }
        Ok(String::from_utf8(vector).unwrap())
    }
    fn read_enum<T, F>(&mut self, _: &str, f: F) -> IoResult<T> where
        F: FnOnce(&mut DecoderReader<'a, R>) -> IoResult<T> {
            f(self)
        }
    fn read_enum_variant<T, F>(&mut self, _: &[&str], f: F) -> IoResult<T> where
        F: FnOnce(&mut DecoderReader<'a, R>, uint) -> IoResult<T> {
            let id = try!(self.read_uint());
            f(self, id)
        }
    fn read_enum_variant_arg<T, F>(&mut self, _: uint, f: F) -> IoResult<T> where
        F: FnOnce(&mut DecoderReader<'a, R>) -> IoResult<T> {
            f(self)
        }
    fn read_enum_struct_variant<T, F>(&mut self, names: &[&str], f: F) -> IoResult<T> where
        F: FnOnce(&mut DecoderReader<'a, R>, uint) -> IoResult<T> {
            self.read_enum_variant(names, f)
        }
    fn read_enum_struct_variant_field<T, F>(&mut self,
                                            _: &str,
                                            f_idx: uint,
                                            f: F)
        -> IoResult<T> where
            F: FnOnce(&mut DecoderReader<'a, R>) -> IoResult<T> {
                self.read_enum_variant_arg(f_idx, f)
            }
    fn read_struct<T, F>(&mut self, _: &str, _: uint, f: F) -> IoResult<T> where
        F: FnOnce(&mut DecoderReader<'a, R>) -> IoResult<T> {
            f(self)
        }
    fn read_struct_field<T, F>(&mut self,
                               _: &str,
                               _: uint,
                               f: F)
        -> IoResult<T> where
            F: FnOnce(&mut DecoderReader<'a, R>) -> IoResult<T> {
                f(self)
            }
    fn read_tuple<T, F>(&mut self, _: uint, f: F) -> IoResult<T> where
        F: FnOnce(&mut DecoderReader<'a, R>) -> IoResult<T> {
            f(self)
        }
    fn read_tuple_arg<T, F>(&mut self, _: uint, f: F) -> IoResult<T> where
        F: FnOnce(&mut DecoderReader<'a, R>) -> IoResult<T> {
            f(self)
        }
    fn read_tuple_struct<T, F>(&mut self, _: &str, len: uint, f: F) -> IoResult<T> where
        F: FnOnce(&mut DecoderReader<'a, R>) -> IoResult<T> {
            self.read_tuple(len, f)
        }
    fn read_tuple_struct_arg<T, F>(&mut self, a_idx: uint, f: F) -> IoResult<T> where
        F: FnOnce(&mut DecoderReader<'a, R>) -> IoResult<T> {
            self.read_tuple_arg(a_idx, f)
        }
    fn read_option<T, F>(&mut self, f: F) -> IoResult<T> where
        F: FnOnce(&mut DecoderReader<'a, R>, bool) -> IoResult<T> {
            match try!(self.reader.read_u8()) {
                1 => f(self, true),
                _ => f(self, false)
            }
        }
    fn read_seq<T, F>(&mut self, f: F) -> IoResult<T> where
        F: FnOnce(&mut DecoderReader<'a, R>, uint) -> IoResult<T> {
            let len = try!(self.read_uint());
            f(self, len)
        }
    fn read_seq_elt<T, F>(&mut self, _: uint, f: F) -> IoResult<T> where
        F: FnOnce(&mut DecoderReader<'a, R>) -> IoResult<T> {
            f(self)
        }
    fn read_map<T, F>(&mut self, f: F) -> IoResult<T> where
        F: FnOnce(&mut DecoderReader<'a, R>, uint) -> IoResult<T> {
            let len = try!(self.read_uint());
            f(self, len)
        }
    fn read_map_elt_key<T, F>(&mut self, _: uint, f: F) -> IoResult<T> where
        F: FnOnce(&mut DecoderReader<'a, R>) -> IoResult<T> {
            f(self)
        }
    fn read_map_elt_val<T, F>(&mut self, _: uint, f: F) -> IoResult<T> where
        F: FnOnce(&mut DecoderReader<'a, R>) -> IoResult<T> {
            f(self)
        }
    fn error(&mut self, err: &str) -> IoError {
        IoError {
            kind: OtherIoError,
            desc: "failure decoding or something, I don't know",
            detail: Some(err.to_string())
        }
    }
}
<commit_msg>Change closures to FnMut to correspond to change in serialize crate.<commit_after>use std::io::{Buffer, Reader, IoError, IoResult, OtherIoError};
use serialize::Decoder;

pub struct DecoderReader<'a, R: 'a> {
    reader: &'a mut R
}

impl<'a, R: Reader+Buffer> DecoderReader<'a, R> {
    pub fn new(r: &'a mut R) -> DecoderReader<'a, R> {
        DecoderReader {
            reader: r
        }
    }
}

impl<'a, R: Reader+Buffer> Decoder<IoError> for DecoderReader<'a, R> {
    fn read_nil(&mut self) -> IoResult<()> {
        Ok(())
    }
    fn read_uint(&mut self) -> IoResult<uint> {
        self.read_u64().map(|x| x as uint)
    }
    fn read_u64(&mut self) -> IoResult<u64> {
        self.reader.read_be_u64()
    }
    fn read_u32(&mut self) -> IoResult<u32> {
        self.reader.read_be_u32()
    }
    fn read_u16(&mut self) -> IoResult<u16> {
        self.reader.read_be_u16()
    }
    fn read_u8(&mut self) -> IoResult<u8> {
        self.reader.read_u8()
    }
    fn read_int(&mut self) -> IoResult<int> {
        self.read_i64().map(|x| x as int)
    }
    fn read_i64(&mut self) -> IoResult<i64> {
        self.reader.read_be_i64()
    }
    fn read_i32(&mut self) -> IoResult<i32> {
        self.reader.read_be_i32()
    }
    fn read_i16(&mut self) -> IoResult<i16> {
        self.reader.read_be_i16()
    }
    fn read_i8(&mut self) -> IoResult<i8> {
        self.reader.read_i8()
    }
    fn read_bool(&mut self) -> IoResult<bool> {
        match try!(self.reader.read_i8()) {
            1 => Ok(true),
            _ => Ok(false)
        }
    }
    fn read_f64(&mut self) -> IoResult<f64> {
        self.reader.read_be_f64()
    }
    fn read_f32(&mut self) -> IoResult<f32> {
        self.reader.read_be_f32()
    }
    fn read_char(&mut self) -> IoResult<char> {
        self.reader.read_char()
    }
    fn read_str(&mut self) -> IoResult<String> {
        let len = try!(self.read_uint());
        let mut vector = Vec::with_capacity(len as uint);
        for _ in range(0, len) {
            vector.push(try!(self.reader.read_u8()));
        }
        Ok(String::from_utf8(vector).unwrap())
    }
    fn read_enum<T, F>(&mut self, _: &str, f: F) -> IoResult<T> where
        F: FnOnce(&mut DecoderReader<'a, R>) -> IoResult<T> {
            f(self)
        }
    fn read_enum_variant<T, F>(&mut self, _: &[&str], mut f: F) -> IoResult<T> where
        F: FnMut(&mut DecoderReader<'a, R>, uint) -> IoResult<T> {
            let id = try!(self.read_uint());
            f(self, id)
        }
    fn read_enum_variant_arg<T, F>(&mut self, _: uint, f: F) -> IoResult<T> where
        F: FnOnce(&mut DecoderReader<'a, R>) -> IoResult<T> {
            f(self)
        }
    fn read_enum_struct_variant<T, F>(&mut self, names: &[&str], f: F) -> IoResult<T> where
        F: FnMut(&mut DecoderReader<'a, R>, uint) -> IoResult<T> {
            self.read_enum_variant(names, f)
        }
    fn read_enum_struct_variant_field<T, F>(&mut self,
                                            _: &str,
                                            f_idx: uint,
                                            f: F)
        -> IoResult<T> where
            F: FnOnce(&mut DecoderReader<'a, R>) -> IoResult<T> {
                self.read_enum_variant_arg(f_idx, f)
            }
    fn read_struct<T, F>(&mut self, _: &str, _: uint, f: F) -> IoResult<T> where
        F: FnOnce(&mut DecoderReader<'a, R>) -> IoResult<T> {
            f(self)
        }
    fn read_struct_field<T, F>(&mut self,
                               _: &str,
                               _: uint,
                               f: F)
        -> IoResult<T> where
            F: FnOnce(&mut DecoderReader<'a, R>) -> IoResult<T> {
                f(self)
            }
    fn read_tuple<T, F>(&mut self, _: uint, f: F) -> IoResult<T> where
        F: FnOnce(&mut DecoderReader<'a, R>) -> IoResult<T> {
            f(self)
        }
    fn read_tuple_arg<T, F>(&mut self, _: uint, f: F) -> IoResult<T> where
        F: FnOnce(&mut DecoderReader<'a, R>) -> IoResult<T> {
            f(self)
        }
    fn read_tuple_struct<T, F>(&mut self, _: &str, len: uint, f: F) -> IoResult<T> where
        F: FnOnce(&mut DecoderReader<'a, R>) -> IoResult<T> {
            self.read_tuple(len, f)
        }
    fn read_tuple_struct_arg<T, F>(&mut self, a_idx: uint, f: F) -> IoResult<T> where
        F: FnOnce(&mut DecoderReader<'a, R>) -> IoResult<T> {
            self.read_tuple_arg(a_idx, f)
        }
    fn read_option<T, F>(&mut self, mut f: F) -> IoResult<T> where
        F: FnMut(&mut DecoderReader<'a, R>, bool) -> IoResult<T> {
            match try!(self.reader.read_u8()) {
                1 => f(self, true),
                _ => f(self, false)
            }
        }
    fn read_seq<T, F>(&mut self, f: F) -> IoResult<T> where
        F: FnOnce(&mut DecoderReader<'a, R>, uint) -> IoResult<T> {
            let len = try!(self.read_uint());
            f(self, len)
        }
    fn read_seq_elt<T, F>(&mut self, _: uint, f: F) -> IoResult<T> where
        F: FnOnce(&mut DecoderReader<'a, R>) -> IoResult<T> {
            f(self)
        }
    fn read_map<T, F>(&mut self, f: F) -> IoResult<T> where
        F: FnOnce(&mut DecoderReader<'a, R>, uint) -> IoResult<T> {
            let len = try!(self.read_uint());
            f(self, len)
        }
    fn read_map_elt_key<T, F>(&mut self, _: uint, f: F) -> IoResult<T> where
        F: FnOnce(&mut DecoderReader<'a, R>) -> IoResult<T> {
            f(self)
        }
    fn read_map_elt_val<T, F>(&mut self, _: uint, f: F) -> IoResult<T> where
        F: FnOnce(&mut DecoderReader<'a, R>) -> IoResult<T> {
            f(self)
        }
    fn error(&mut self, err: &str) -> IoError {
        IoError {
            kind: OtherIoError,
            desc: "failure decoding or something, I don't know",
            detail: Some(err.to_string())
        }
    }
}
<|endoftext|>