AkurAI-Framework / crates/sql/src/lexer.rs

Latest change 5a1781679346dc042a413c85c2ddfc034f219ae0 - feat(sql): lexer + parser for a SQL subset (AST) by Ólafur Búi Ólafsson

//! Turn SQL text into a stream of position-tagged [`SpannedToken`]s.
//!
//! Operates over a `Vec<char>` for simple, correct indexing (positions are
//! 0-based character offsets). Whitespace is skipped; every other run of input
//! becomes exactly one token or a positioned [`SqlError`].

use crate::error::SqlError;
use crate::token::{Keyword, SpannedToken, Token};

/// Tokenize a full SQL string. Returns every token in order (no trailing EOF
/// marker — the parser checks for end-of-input via the slice length).
pub fn tokenize(input: &str) -> Result<Vec<SpannedToken>, SqlError> {
    let mut lexer = Lexer {
        chars: input.chars().collect(),
        pos: 0,
    };
    lexer.run()
}

struct Lexer {
    chars: Vec<char>,
    pos: usize,
}

impl Lexer {
    fn peek(&self) -> Option<char> {
        self.chars.get(self.pos).copied()
    }

    fn peek_at(&self, offset: usize) -> Option<char> {
        self.chars.get(self.pos + offset).copied()
    }

    fn run(&mut self) -> Result<Vec<SpannedToken>, SqlError> {
        let mut tokens = Vec::new();
        loop {
            self.skip_ws();
            let start = self.pos;
            let c = match self.peek() {
                Some(c) => c,
                None => return Ok(tokens),
            };
            let token = match c {
                '*' => self.single(Token::Star),
                ',' => self.single(Token::Comma),
                '(' => self.single(Token::LParen),
                ')' => self.single(Token::RParen),
                ';' => self.single(Token::Semicolon),
                '=' => self.single(Token::Eq),
                '<' => self.lex_lt(),
                '>' => self.lex_gt(),
                '!' => self.lex_bang()?,
                '\'' => self.lex_string()?,
                c if c.is_ascii_digit() => self.lex_number()?,
                c if is_ident_start(c) => self.lex_word(),
                c => return Err(SqlError::new(format!("unexpected character '{c}'"), start)),
            };
            tokens.push(SpannedToken { token, pos: start });
        }
    }

    /// Consume one char and emit the given token.
    fn single(&mut self, token: Token) -> Token {
        self.pos += 1;
        token
    }

    fn skip_ws(&mut self) {
        while matches!(self.peek(), Some(c) if c.is_whitespace()) {
            self.pos += 1;
        }
    }

    fn lex_lt(&mut self) -> Token {
        self.pos += 1; // '<'
        match self.peek() {
            Some('=') => self.single(Token::LtEq),
            Some('>') => self.single(Token::NotEq), // SQL's `<>`
            _ => Token::Lt,
        }
    }

    fn lex_gt(&mut self) -> Token {
        self.pos += 1; // '>'
        match self.peek() {
            Some('=') => self.single(Token::GtEq),
            _ => Token::Gt,
        }
    }

    fn lex_bang(&mut self) -> Result<Token, SqlError> {
        let start = self.pos;
        self.pos += 1; // '!'
        match self.peek() {
            Some('=') => Ok(self.single(Token::NotEq)),
            _ => Err(SqlError::new("expected '=' after '!'", start)),
        }
    }

    /// A single-quoted string literal. Two consecutive quotes (`''`) are an
    /// escaped single quote inside the string.
    fn lex_string(&mut self) -> Result<Token, SqlError> {
        let open = self.pos;
        self.pos += 1; // opening quote
        let mut s = String::new();
        loop {
            match self.peek() {
                None => return Err(SqlError::new("unterminated string literal", open)),
                Some('\'') => {
                    if self.peek_at(1) == Some('\'') {
                        s.push('\'');
                        self.pos += 2;
                    } else {
                        self.pos += 1; // closing quote
                        return Ok(Token::Str(s));
                    }
                }
                Some(c) => {
                    s.push(c);
                    self.pos += 1;
                }
            }
        }
    }

    /// An integer or float literal. The leading digit is already confirmed.
    fn lex_number(&mut self) -> Result<Token, SqlError> {
        let start = self.pos;
        let mut is_float = false;
        while matches!(self.peek(), Some(c) if c.is_ascii_digit()) {
            self.pos += 1;
        }
        if self.peek() == Some('.') {
            is_float = true;
            self.pos += 1;
            while matches!(self.peek(), Some(c) if c.is_ascii_digit()) {
                self.pos += 1;
            }
        }
        if matches!(self.peek(), Some('e' | 'E')) {
            is_float = true;
            self.pos += 1;
            if matches!(self.peek(), Some('+' | '-')) {
                self.pos += 1;
            }
            while matches!(self.peek(), Some(c) if c.is_ascii_digit()) {
                self.pos += 1;
            }
        }
        let text: String = self.chars[start..self.pos].iter().collect();
        if is_float {
            text.parse::<f64>()
                .map(Token::Float)
                .map_err(|_| SqlError::new("invalid number literal", start))
        } else {
            text.parse::<i64>()
                .map(Token::Int)
                .map_err(|_| SqlError::new("integer literal out of range", start))
        }
    }

    /// An identifier or a keyword. Keywords match case-insensitively.
    fn lex_word(&mut self) -> Token {
        let start = self.pos;
        while matches!(self.peek(), Some(c) if is_ident_continue(c)) {
            self.pos += 1;
        }
        let word: String = self.chars[start..self.pos].iter().collect();
        match Keyword::from_upper(&word.to_ascii_uppercase()) {
            Some(kw) => Token::Keyword(kw),
            None => Token::Ident(word),
        }
    }
}

fn is_ident_start(c: char) -> bool {
    c == '_' || c.is_ascii_alphabetic()
}

fn is_ident_continue(c: char) -> bool {
    c == '_' || c.is_ascii_alphanumeric()
}

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

    fn toks(sql: &str) -> Vec<Token> {
        tokenize(sql)
            .unwrap()
            .into_iter()
            .map(|t| t.token)
            .collect()
    }

    #[test]
    fn lexes_keywords_case_insensitively() {
        assert_eq!(
            toks("select FROM wHeRe"),
            vec![
                Token::Keyword(Keyword::Select),
                Token::Keyword(Keyword::From),
                Token::Keyword(Keyword::Where),
            ]
        );
    }

    #[test]
    fn lexes_operators_and_punctuation() {
        assert_eq!(
            toks("= != <> < <= > >= * , ( ) ;"),
            vec![
                Token::Eq,
                Token::NotEq,
                Token::NotEq,
                Token::Lt,
                Token::LtEq,
                Token::Gt,
                Token::GtEq,
                Token::Star,
                Token::Comma,
                Token::LParen,
                Token::RParen,
                Token::Semicolon,
            ]
        );
    }

    #[test]
    fn lexes_identifiers_and_numbers() {
        assert_eq!(
            toks("users id_2 42 2.5 1e3"),
            vec![
                Token::Ident("users".into()),
                Token::Ident("id_2".into()),
                Token::Int(42),
                Token::Float(2.5),
                Token::Float(1000.0),
            ]
        );
    }

    #[test]
    fn lexes_string_with_escaped_quote() {
        assert_eq!(toks("'it''s ok'"), vec![Token::Str("it's ok".into())]);
        assert_eq!(toks("''"), vec![Token::Str(String::new())]);
    }

    #[test]
    fn tracks_positions() {
        let spanned = tokenize("a = 1").unwrap();
        assert_eq!(spanned[0].pos, 0);
        assert_eq!(spanned[1].pos, 2);
        assert_eq!(spanned[2].pos, 4);
    }

    #[test]
    fn errors_are_positioned_not_panics() {
        assert_eq!(tokenize("a # b").unwrap_err().position, 2);
        assert_eq!(tokenize("'open").unwrap_err().position, 0);
        assert!(tokenize("a ! b").is_err());
    }
}