Skip to main content

vecdb_core/planner/sql/
lexer.rs

1use crate::errors::{Result, VecDbError};
2use crate::types::Vector;
3
4// ─────────────────────────────────────────────────────────────────
5// Token
6// ─────────────────────────────────────────────────────────────────
7
8#[derive(Debug, Clone, PartialEq)]
9pub enum Token {
10    // ── Keywords ─────────────────────────────────────────────────
11    Select,
12    From,
13    Where,
14    And,
15    Or,
16    Not,
17    Order,
18    By,
19    Asc,
20    Desc,
21    Limit,
22    Like,
23    True,
24    False,
25    // ── Functions ─────────────────────────────────────────────────
26    VectorSim,
27    Embed,
28    // ── Identifiers and literals ──────────────────────────────────
29    Identifier(String),
30    StringLit(String),
31    IntLit(i64),
32    FloatLit(f64),
33    /// Inline vector literal `[f32, ...]`.
34    VectorLit(Vector),
35    // ── Comparison operators ──────────────────────────────────────
36    Eq,    // =
37    NotEq, // != or <>
38    Lt,    // <
39    LtEq,  // <=
40    Gt,    // >
41    GtEq,  // >=
42    // ── JSON path operators ───────────────────────────────────────
43    Arrow,     // ->
44    ArrowText, // ->>
45    // ── Punctuation ──────────────────────────────────────────────
46    LParen,    // (
47    RParen,    // )
48    LBracket,  // [
49    RBracket,  // ]
50    Comma,     // ,
51    Dot,       // .
52    Star,      // *
53    Semicolon, // ;
54    Eof,
55}
56
57// ─────────────────────────────────────────────────────────────────
58// Lexer
59// ─────────────────────────────────────────────────────────────────
60
61pub struct Lexer {
62    input: Vec<char>,
63    pos: usize,
64}
65
66impl Lexer {
67    pub fn new(input: &str) -> Self {
68        Self {
69            input: input.chars().collect(),
70            pos: 0,
71        }
72    }
73
74    /// Tokenize the entire input.
75    pub fn tokenize(&mut self) -> Result<Vec<Token>> {
76        let mut tokens = Vec::new();
77        loop {
78            let tok = self.next_token()?;
79            let done = tok == Token::Eof;
80            tokens.push(tok);
81            if done {
82                break;
83            }
84        }
85        Ok(tokens)
86    }
87
88    // ── Low-level helpers ─────────────────────────────────────────
89
90    fn peek(&self) -> Option<char> {
91        self.input.get(self.pos).copied()
92    }
93
94    fn peek_next(&self) -> Option<char> {
95        self.input.get(self.pos + 1).copied()
96    }
97
98    fn advance(&mut self) -> Option<char> {
99        let ch = self.input.get(self.pos).copied();
100        if ch.is_some() {
101            self.pos += 1;
102        }
103        ch
104    }
105
106    fn skip_whitespace(&mut self) {
107        while matches!(self.peek(), Some(c) if c.is_whitespace()) {
108            self.advance();
109        }
110    }
111
112    // ── Main dispatch ─────────────────────────────────────────────
113
114    fn next_token(&mut self) -> Result<Token> {
115        self.skip_whitespace();
116
117        match self.peek() {
118            None => Ok(Token::Eof),
119            Some(c) => match c {
120                '(' => {
121                    self.advance();
122                    Ok(Token::LParen)
123                }
124                ')' => {
125                    self.advance();
126                    Ok(Token::RParen)
127                }
128                '[' => {
129                    self.advance();
130                    self.scan_vector_lit()
131                }
132                ']' => {
133                    self.advance();
134                    Ok(Token::RBracket)
135                }
136                ',' => {
137                    self.advance();
138                    Ok(Token::Comma)
139                }
140                '.' => {
141                    self.advance();
142                    Ok(Token::Dot)
143                }
144                '*' => {
145                    self.advance();
146                    Ok(Token::Star)
147                }
148                ';' => {
149                    self.advance();
150                    Ok(Token::Semicolon)
151                }
152                '=' => {
153                    self.advance();
154                    Ok(Token::Eq)
155                }
156                '!' => {
157                    self.advance();
158                    if self.peek() == Some('=') {
159                        self.advance();
160                        Ok(Token::NotEq)
161                    } else {
162                        Err(VecDbError::InvalidQuery(format!(
163                            "unexpected '!' at pos {}",
164                            self.pos
165                        )))
166                    }
167                }
168                '<' => {
169                    self.advance();
170                    match self.peek() {
171                        Some('=') => {
172                            self.advance();
173                            Ok(Token::LtEq)
174                        }
175                        Some('>') => {
176                            self.advance();
177                            Ok(Token::NotEq)
178                        }
179                        _ => Ok(Token::Lt),
180                    }
181                }
182                '>' => {
183                    self.advance();
184                    if self.peek() == Some('=') {
185                        self.advance();
186                        Ok(Token::GtEq)
187                    } else {
188                        Ok(Token::Gt)
189                    }
190                }
191                '-' => {
192                    self.advance();
193                    match self.peek() {
194                        Some('>') => {
195                            self.advance();
196                            if self.peek() == Some('>') {
197                                self.advance();
198                                Ok(Token::ArrowText) // ->>
199                            } else {
200                                Ok(Token::Arrow) // ->
201                            }
202                        }
203                        Some(d) if d.is_ascii_digit() => self.scan_number(true),
204                        _ => Err(VecDbError::InvalidQuery(format!(
205                            "unexpected '-' at pos {}",
206                            self.pos
207                        ))),
208                    }
209                }
210                '\'' => self.scan_string(),
211                c if c.is_ascii_digit() => self.scan_number(false),
212                c if c.is_alphabetic() || c == '_' => self.scan_ident_or_keyword(),
213                other => Err(VecDbError::InvalidQuery(format!(
214                    "unexpected character '{}' at pos {}",
215                    other, self.pos
216                ))),
217            },
218        }
219    }
220
221    // ── Single-quoted string ──────────────────────────────────────
222
223    /// Scan `'...'` — `''` inside the string is an escaped single-quote.
224    fn scan_string(&mut self) -> Result<Token> {
225        self.advance(); // consume opening '
226        let mut s = String::new();
227        loop {
228            match self.advance() {
229                None => {
230                    return Err(VecDbError::InvalidQuery(
231                        "unterminated string literal".into(),
232                    ))
233                }
234                Some('\'') => {
235                    if self.peek() == Some('\'') {
236                        self.advance(); // consume second '
237                        s.push('\'');
238                    } else {
239                        break;
240                    }
241                }
242                Some(c) => s.push(c),
243            }
244        }
245        Ok(Token::StringLit(s))
246    }
247
248    // ── Number ────────────────────────────────────────────────────
249
250    /// Scan an integer or float.  `negative` is true when the `-` has
251    /// already been consumed by the caller.
252    fn scan_number(&mut self, negative: bool) -> Result<Token> {
253        let mut s = String::new();
254        if negative {
255            s.push('-');
256        }
257        while matches!(self.peek(), Some(c) if c.is_ascii_digit()) {
258            s.push(self.advance().unwrap());
259        }
260        // Detect float: digit `.` digit ...
261        let is_float =
262            self.peek() == Some('.') && matches!(self.peek_next(), Some(c) if c.is_ascii_digit());
263        if is_float {
264            s.push(self.advance().unwrap()); // '.'
265            while matches!(self.peek(), Some(c) if c.is_ascii_digit()) {
266                s.push(self.advance().unwrap());
267            }
268            // Optional exponent `e` / `E`
269            if matches!(self.peek(), Some('e') | Some('E')) {
270                s.push(self.advance().unwrap());
271                if matches!(self.peek(), Some('+') | Some('-')) {
272                    s.push(self.advance().unwrap());
273                }
274                while matches!(self.peek(), Some(c) if c.is_ascii_digit()) {
275                    s.push(self.advance().unwrap());
276                }
277            }
278            let v: f64 = s
279                .parse()
280                .map_err(|_| VecDbError::InvalidQuery(format!("invalid float literal: {s}")))?;
281            Ok(Token::FloatLit(v))
282        } else {
283            let v: i64 = s
284                .parse()
285                .map_err(|_| VecDbError::InvalidQuery(format!("invalid integer literal: {s}")))?;
286            Ok(Token::IntLit(v))
287        }
288    }
289
290    // ── Inline vector literal  `[f32, ...]` ──────────────────────
291
292    /// Scan the body of `[...]`.  The opening `[` has already been consumed.
293    fn scan_vector_lit(&mut self) -> Result<Token> {
294        let mut floats: Vec<f32> = Vec::new();
295        self.skip_whitespace();
296
297        // Empty vector `[]`
298        if self.peek() == Some(']') {
299            self.advance();
300            return Ok(Token::VectorLit(floats));
301        }
302
303        loop {
304            self.skip_whitespace();
305
306            let negative = if self.peek() == Some('-') {
307                self.advance();
308                true
309            } else {
310                false
311            };
312
313            let mut s = String::new();
314            if negative {
315                s.push('-');
316            }
317            while matches!(self.peek(), Some(c) if c.is_ascii_digit()) {
318                s.push(self.advance().unwrap());
319            }
320            if self.peek() == Some('.') {
321                s.push(self.advance().unwrap());
322                while matches!(self.peek(), Some(c) if c.is_ascii_digit()) {
323                    s.push(self.advance().unwrap());
324                }
325            }
326            // optional exponent
327            if matches!(self.peek(), Some('e') | Some('E')) {
328                s.push(self.advance().unwrap());
329                if matches!(self.peek(), Some('+') | Some('-')) {
330                    s.push(self.advance().unwrap());
331                }
332                while matches!(self.peek(), Some(c) if c.is_ascii_digit()) {
333                    s.push(self.advance().unwrap());
334                }
335            }
336            if s.is_empty() || s == "-" {
337                return Err(VecDbError::InvalidQuery(
338                    "expected number in vector literal".into(),
339                ));
340            }
341            let v: f32 = s.parse().map_err(|_| {
342                VecDbError::InvalidQuery(format!("invalid float in vector literal: {s}"))
343            })?;
344            floats.push(v);
345
346            self.skip_whitespace();
347            match self.peek() {
348                Some(',') => {
349                    self.advance();
350                }
351                Some(']') => {
352                    self.advance();
353                    break;
354                }
355                other => {
356                    return Err(VecDbError::InvalidQuery(format!(
357                        "expected ',' or ']' in vector literal, got {:?}",
358                        other
359                    )))
360                }
361            }
362        }
363        Ok(Token::VectorLit(floats))
364    }
365
366    // ── Identifier / keyword ──────────────────────────────────────
367
368    fn scan_ident_or_keyword(&mut self) -> Result<Token> {
369        let mut s = String::new();
370        while matches!(self.peek(), Some(c) if c.is_alphanumeric() || c == '_') {
371            s.push(self.advance().unwrap());
372        }
373        Ok(match s.to_uppercase().as_str() {
374            "SELECT" => Token::Select,
375            "FROM" => Token::From,
376            "WHERE" => Token::Where,
377            "AND" => Token::And,
378            "OR" => Token::Or,
379            "NOT" => Token::Not,
380            "ORDER" => Token::Order,
381            "BY" => Token::By,
382            "ASC" => Token::Asc,
383            "DESC" => Token::Desc,
384            "LIMIT" => Token::Limit,
385            "LIKE" => Token::Like,
386            "TRUE" => Token::True,
387            "FALSE" => Token::False,
388            "VECTOR_SIM" => Token::VectorSim,
389            "EMBED" => Token::Embed,
390            _ => Token::Identifier(s),
391        })
392    }
393}