kappe/src/parser.rs

use std::{iter::Peekable, slice::Iter};

use crate::ast::{Equals, Field, KIdentifier, KNumber, KString, KValue, Select, Where};
use crate::error::KappeError;
use crate::token::{Token, TokenType};

pub fn parse(tokens: Vec<Token>) -> Result<Select, KappeError> {
    parse_select_expression(&mut tokens.iter().peekable())
}

fn parse_select_expression(tokens: &mut Peekable<Iter<Token>>) -> Result<Select, KappeError> {
    eat(tokens, TokenType::Select)?;
    let fields = parse_fields(tokens)?;
    eat(tokens, TokenType::From)?;
    let source = eat(tokens, TokenType::Identfiier).and_then(|t| Ok(KIdentifier::new(&t.value)))?;

    let where_clause = parse_where_clause(tokens)?;

    Ok(Select {
        fields,
        source,
        where_clause,
    })
}

fn parse_fields(tokens: &mut Peekable<Iter<Token>>) -> Result<Vec<Field>, KappeError> {
    if current_token_is(tokens, TokenType::Star) {
        eat(tokens, TokenType::Star).and_then(|_| Ok(vec![Field::Star]))
    } else {
        let mut fields: Vec<Field> = vec![];

        loop {
            let field = eat(tokens, TokenType::Identfiier)?;
            fields.push(Field::Named(field.value.clone()));

            if current_token_is(tokens, TokenType::Comma) {
                eat(tokens, TokenType::Comma)?;
            } else {
                break;
            }
        }

        Ok(fields)
    }
}

fn parse_where_clause(tokens: &mut Peekable<Iter<Token>>) -> Result<Option<Where>, KappeError> {
    let mut predicates: Vec<Equals> = vec![];
    let where_clause = if current_token_is(tokens, TokenType::Where) {
        eat(tokens, TokenType::Where)?;

        let predicate = parse_predicate(tokens)?;
        predicates.push(predicate);

        if current_token_is(tokens, TokenType::Comma) {
            while current_token_is(tokens, TokenType::Comma) {
                eat(tokens, TokenType::Comma)?;
                let predicate = parse_predicate(tokens)?;
                predicates.push(predicate);
            }
        }

        Some(Where::new(predicates))
    } else {
        None
    };

    Ok(where_clause)
}

fn parse_predicate(tokens: &mut Peekable<Iter<Token>>) -> Result<Equals, KappeError> {
    let field = eat(tokens, TokenType::Identfiier)?;
    let field = KIdentifier::new(&field.value);

    eat(tokens, TokenType::Equals)?;

    let value: KValue = if current_token_is(tokens, TokenType::String) {
        let token = eat(tokens, TokenType::String)?;
        Ok(KValue::String(KString::new(&token.value)))
    } else if current_token_is(tokens, TokenType::Number) {
        let token = eat(tokens, TokenType::Number)?;
        match token.value.parse::<f64>() {
            Ok(n) => Ok(KValue::Number(KNumber::new(n))),
            Err(_) => Err(KappeError::new("Error converting number to f64")),
        }
    } else {
        Err(KappeError::new(
            "Invalid token on right-hand side of equals sign",
        ))
    }?;

    Ok(Equals::new(field, value))
}

fn current_token_is(tokens: &mut Peekable<Iter<Token>>, token_type: TokenType) -> bool {
    match tokens.peek() {
        Some(t) => t.token_type == token_type,
        None => false,
    }
}

fn eat<'a>(
    tokens: &'a mut Peekable<Iter<Token>>,
    token_type: TokenType,
) -> Result<&'a Token, KappeError> {
    match tokens.next() {
        Some(token) => {
            if token.token_type == token_type {
                Ok(token)
            } else {
                Err(KappeError::new(&format!(
                    "Expected {} but got {}",
                    token_type, token.token_type
                )))
            }
        }
        None => Err(KappeError::new("Unexpected end of tokens")),
    }
}

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

    fn _parse(input: &str) -> Result<Select, KappeError> {
        let tokens = scan(input).unwrap();
        parse(tokens)
    }

    #[test]
    fn it_parses_a_simple_select() {
        assert_eq!(
            _parse("SELECT * FROM index").unwrap(),
            Select {
                fields: vec![Field::Star],
                source: KIdentifier::new("index"),
                where_clause: None
            },
        )
    }

    #[test]
    fn it_parses_a_select_with_field() {
        assert_eq!(
            _parse("SELECT field FROM index").unwrap(),
            Select {
                fields: vec![Field::Named("field".to_string())],
                source: KIdentifier::new("index"),
                where_clause: None
            },
        )
    }

    #[test]
    fn it_parses_a_select_with_multiple_fields() {
        assert_eq!(
            _parse("SELECT field_one, field_two FROM index").unwrap(),
            Select {
                fields: vec![
                    Field::Named("field_one".to_string()),
                    Field::Named("field_two".to_string())
                ],
                source: KIdentifier::new("index"),
                where_clause: None
            },
        )
    }

    #[test]
    fn it_parses_a_select_with_a_single_predicate_where_clause() {
        assert_eq!(
            _parse("SELECT * FROM index WHERE field = \"hello\"").unwrap(),
            Select {
                fields: vec![Field::Star],
                source: KIdentifier::new("index"),
                where_clause: Some(Where::new(vec![Equals::new(
                    KIdentifier::new("field"),
                    KValue::String(KString::new("hello")),
                )]))
            },
        )
    }

    #[test]
    fn it_parses_a_select_with_a_numeric_predicate() {
        assert_eq!(
            _parse("SELECT * FROM index WHERE field = 123").unwrap(),
            Select {
                fields: vec![Field::Star],
                source: KIdentifier::new("index"),
                where_clause: Some(Where::new(vec![Equals::new(
                    KIdentifier::new("field"),
                    KValue::Number(KNumber::new(123.0))
                ),]))
            },
        )
    }

    #[test]
    fn it_parses_a_select_with_a_multiple_predicate_where_clause() {
        assert_eq!(
            _parse("SELECT * FROM index WHERE field_one = \"hello\", field_two = \"world\"")
                .unwrap(),
            Select {
                fields: vec![Field::Star],
                source: KIdentifier::new("index"),
                where_clause: Some(Where::new(vec![
                    Equals::new(
                        KIdentifier::new("field_one"),
                        KValue::String(KString::new("hello")),
                    ),
                    Equals::new(
                        KIdentifier::new("field_two"),
                        KValue::String(KString::new("world")),
                    )
                ]))
            },
        )
    }
}