lance_tokenizer/
code_tokenizer.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::{iter::Peekable, str::CharIndices};
5
6use crate::{Token, TokenStream, Tokenizer};
7
8/// Tokenizer for code-like text.
9///
10/// Identifiers are Unicode alphanumeric characters plus `_`. Other characters
11/// are lexical boundaries. When operator indexing is enabled, recognized
12/// multi-character operators use longest-match tokenization and remaining
13/// operator characters are emitted individually.
14///
15/// # Examples
16///
17/// ```
18/// use lance_tokenizer::{CodeLexTokenizer, TextAnalyzer, TokenStream};
19///
20/// let mut analyzer = TextAnalyzer::builder(CodeLexTokenizer::new(true)).build();
21/// let mut stream = analyzer.token_stream("a::b");
22///
23/// assert!(stream.advance());
24/// assert_eq!(stream.token().text, "a");
25/// assert!(stream.advance());
26/// assert_eq!(stream.token().text, "::");
27/// ```
28#[derive(Clone, Default)]
29pub struct CodeLexTokenizer {
30    index_operators: bool,
31    token: Token,
32}
33
34impl CodeLexTokenizer {
35    pub fn new(index_operators: bool) -> Self {
36        Self {
37            index_operators,
38            token: Token::default(),
39        }
40    }
41}
42
43/// Token stream produced by [`CodeLexTokenizer`].
44pub struct CodeLexTokenStream<'a> {
45    text: &'a str,
46    chars: Peekable<CharIndices<'a>>,
47    token: &'a mut Token,
48    index_operators: bool,
49}
50
51impl Tokenizer for CodeLexTokenizer {
52    type TokenStream<'a> = CodeLexTokenStream<'a>;
53
54    fn token_stream<'a>(&'a mut self, text: &'a str) -> Self::TokenStream<'a> {
55        self.token.reset();
56        CodeLexTokenStream {
57            text,
58            chars: text.char_indices().peekable(),
59            token: &mut self.token,
60            index_operators: self.index_operators,
61        }
62    }
63}
64
65fn is_identifier_char(ch: char) -> bool {
66    ch == '_' || ch.is_alphanumeric()
67}
68
69fn is_operator_char(ch: char) -> bool {
70    matches!(
71        ch,
72        '!' | '%' | '&' | '*' | '+' | '-' | '/' | ':' | '<' | '=' | '>' | '?' | '^' | '|' | '~'
73    )
74}
75
76const MULTI_CHAR_OPERATORS: &[&str] = &[
77    ">>>=", "<<=", ">>=", "&&=", "||=", "??=", "**=", "//=", "===", "!==", ">>>", "<=>", "::",
78    "->", "=>", "==", "!=", "<=", ">=", "&&", "||", "++", "--", "+=", "-=", "*=", "/=", "%=", "&=",
79    "|=", "^=", "<<", ">>", "**", "//", "??", ":=", "<-", "|>", "~=",
80];
81
82impl CodeLexTokenStream<'_> {
83    fn search_token_end(&mut self, predicate: impl Fn(char) -> bool) -> usize {
84        while let Some((_, ch)) = self.chars.peek() {
85            if !predicate(*ch) {
86                break;
87            }
88            self.chars.next();
89        }
90        self.chars
91            .peek()
92            .map(|(offset, _)| *offset)
93            .unwrap_or(self.text.len())
94    }
95
96    fn operator_token_end(&mut self, offset_from: usize) -> usize {
97        let remaining = &self.text[offset_from..];
98        let operator_len = MULTI_CHAR_OPERATORS
99            .iter()
100            .filter(|operator| remaining.starts_with(**operator))
101            .map(|operator| operator.len())
102            .max()
103            .unwrap_or(1);
104        let token_end = offset_from + operator_len;
105        while self
106            .chars
107            .peek()
108            .is_some_and(|(offset, _)| *offset < token_end)
109        {
110            self.chars.next();
111        }
112        token_end
113    }
114}
115
116impl TokenStream for CodeLexTokenStream<'_> {
117    fn advance(&mut self) -> bool {
118        self.token.text.clear();
119        while let Some((offset_from, ch)) = self.chars.next() {
120            let token_end = if is_identifier_char(ch) {
121                self.search_token_end(is_identifier_char)
122            } else if self.index_operators && is_operator_char(ch) {
123                self.operator_token_end(offset_from)
124            } else {
125                continue;
126            };
127
128            self.token.position = self.token.position.wrapping_add(1);
129            self.token.position_length = 1;
130            self.token.offset_from = offset_from;
131            self.token.offset_to = token_end;
132            self.token.text.push_str(&self.text[offset_from..token_end]);
133            return true;
134        }
135        false
136    }
137
138    fn token(&self) -> &Token {
139        self.token
140    }
141
142    fn token_mut(&mut self) -> &mut Token {
143        self.token
144    }
145}
146
147#[cfg(test)]
148mod tests {
149    use crate::{CodeLexTokenizer, TextAnalyzer, Token};
150
151    fn collect_tokens(text: &str, index_operators: bool) -> Vec<Token> {
152        let mut analyzer = TextAnalyzer::builder(CodeLexTokenizer::new(index_operators)).build();
153        let mut stream = analyzer.token_stream(text);
154        let mut tokens = Vec::new();
155        stream.process(&mut |token| tokens.push(token.clone()));
156        tokens
157    }
158
159    #[test]
160    fn test_code_lex_tokenizer_identifiers() {
161        let tokens = collect_tokens("std::vector user-name parse.HTML2JSON", false);
162        let texts = tokens
163            .iter()
164            .map(|token| token.text.as_str())
165            .collect::<Vec<_>>();
166        assert_eq!(
167            texts,
168            vec!["std", "vector", "user", "name", "parse", "HTML2JSON"]
169        );
170    }
171
172    #[test]
173    fn test_code_lex_tokenizer_operators() {
174        let tokens = collect_tokens("a::b != c->d", true);
175        let texts = tokens
176            .iter()
177            .map(|token| token.text.as_str())
178            .collect::<Vec<_>>();
179        assert_eq!(texts, vec!["a", "::", "b", "!=", "c", "->", "d"]);
180    }
181
182    #[test]
183    fn test_code_lex_tokenizer_splits_adjacent_operators() {
184        let tokens = collect_tokens("value.parse::<usize>()", true);
185        let texts = tokens
186            .iter()
187            .map(|token| token.text.as_str())
188            .collect::<Vec<_>>();
189        assert_eq!(texts, vec!["value", "parse", "::", "<", "usize", ">"]);
190    }
191}
lance_tokenizer/code_tokenizer.rs

lance_tokenizer/
code_tokenizer.rs
