ailloli_ui_editor/code/
syntax.rs

1//! Deterministic Rust syntax tokenization and optional Tree-sitter enrichment.
2
3use std::ops::Range;
4
5/// Syntax token emitted by code highlighters.
6///
7/// The half-open range uses UTF-8 byte offsets into the exact source passed to
8/// the highlighter.
9///
10/// # Examples
11///
12/// ```
13/// use ailloli_ui_editor::code::{SyntaxKind, SyntaxToken};
14/// let token = SyntaxToken { range: 0..2, kind: SyntaxKind::Keyword };
15/// assert_eq!(token.range, 0..2);
16/// ```
17#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
18pub struct SyntaxToken {
19    /// Half-open UTF-8 source byte range.
20    pub range: Range<usize>,
21    /// Language-neutral semantic paint category.
22    pub kind: SyntaxKind,
23}
24
25/// Language-neutral syntax token categories.
26///
27/// # Examples
28///
29/// ```
30/// use ailloli_ui_editor::code::SyntaxKind;
31/// assert_ne!(SyntaxKind::Comment, SyntaxKind::Identifier);
32/// ```
33#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
34pub enum SyntaxKind {
35    /// Reserved language word or structural keyword-like node.
36    Keyword,
37    /// Type name or primitive type.
38    Type,
39    /// Function or macro name.
40    Function,
41    /// String or character literal.
42    String,
43    /// Numeric literal.
44    Number,
45    /// Line or block comment.
46    Comment,
47    /// Operator sequence.
48    Operator,
49    /// Delimiter or punctuation byte.
50    Punctuation,
51    /// Identifier or Rust lifetime.
52    Identifier,
53}
54
55/// ASCII Rust keywords recognized by the lexical highlighter.
56const RUST_KEYWORDS: &[&str] = &[
57    "as", "async", "await", "break", "const", "continue", "crate", "dyn", "else", "enum", "extern",
58    "false", "fn", "for", "if", "impl", "in", "let", "loop", "match", "mod", "move", "mut", "pub",
59    "ref", "return", "self", "Self", "static", "struct", "super", "trait", "true", "type",
60    "unsafe", "use", "where", "while",
61];
62
63/// Tokenizes a useful ASCII-oriented subset of Rust without parsing.
64///
65/// The deterministic, linear scan recognizes line comments, quoted strings and
66/// chars, raw strings, lifetimes, numeric/identifier runs, keywords, delimiters,
67/// and operators. It then marks an identifier following `fn` and identifiers
68/// followed by optional ASCII whitespace plus `!` as functions. Non-ASCII
69/// identifier bytes and unrecognized punctuation are skipped. Unterminated
70/// strings/comments extend to EOF; returned ranges are ordered UTF-8 bytes.
71///
72/// # Examples
73///
74/// ```
75/// use ailloli_ui_editor::code::{highlight_rust_lexical, SyntaxKind};
76/// let text = "fn main() { println!(\"hi\"); }";
77/// let tokens = highlight_rust_lexical(text);
78/// assert!(tokens.iter().any(|t| t.kind == SyntaxKind::Keyword && &text[t.range.clone()] == "fn"));
79/// assert!(tokens.iter().any(|t| t.kind == SyntaxKind::Function && &text[t.range.clone()] == "main"));
80/// assert!(tokens.iter().any(|t| t.kind == SyntaxKind::String && &text[t.range.clone()] == "\"hi\""));
81/// ```
82pub fn highlight_rust_lexical(text: &str) -> Vec<SyntaxToken> {
83    let bytes = text.as_bytes();
84    let mut tokens = Vec::new();
85    let mut i = 0;
86    while i < bytes.len() {
87        match bytes[i] {
88            b'/' if bytes.get(i + 1) == Some(&b'/') => {
89                let start = i;
90                i += 2;
91                while i < bytes.len() && bytes[i] != b'\n' {
92                    i += 1;
93                }
94                tokens.push(token(start..i, SyntaxKind::Comment));
95            }
96            b'"' => {
97                let start = i;
98                i += 1;
99                while i < bytes.len() {
100                    if bytes[i] == b'\\' {
101                        i = (i + 2).min(bytes.len());
102                    } else if bytes[i] == b'"' {
103                        i += 1;
104                        break;
105                    } else {
106                        i += 1;
107                    }
108                }
109                tokens.push(token(start..i, SyntaxKind::String));
110            }
111            b'r' if raw_string_end(text, i).is_some() => {
112                let start = i;
113                i = raw_string_end(text, i).unwrap_or(i + 1);
114                tokens.push(token(start..i, SyntaxKind::String));
115            }
116            b'\'' if lifetime_end(bytes, i).is_some() => {
117                let start = i;
118                i = lifetime_end(bytes, i).unwrap_or(i + 1);
119                tokens.push(token(start..i, SyntaxKind::Identifier));
120            }
121            b'\'' => {
122                let start = i;
123                i += 1;
124                while i < bytes.len() {
125                    if bytes[i] == b'\\' {
126                        i = (i + 2).min(bytes.len());
127                    } else if bytes[i] == b'\'' {
128                        i += 1;
129                        break;
130                    } else {
131                        i += 1;
132                    }
133                }
134                tokens.push(token(start..i, SyntaxKind::String));
135            }
136            b'0'..=b'9' => {
137                let start = i;
138                i += 1;
139                while i < bytes.len() && is_ident_continue(bytes[i]) {
140                    i += 1;
141                }
142                tokens.push(token(start..i, SyntaxKind::Number));
143            }
144            ch if is_ident_start(ch) => {
145                let start = i;
146                i += 1;
147                while i < bytes.len() && is_ident_continue(bytes[i]) {
148                    i += 1;
149                }
150                let word = &text[start..i];
151                let kind = if RUST_KEYWORDS.contains(&word) {
152                    SyntaxKind::Keyword
153                } else if word.chars().next().is_some_and(char::is_uppercase) {
154                    SyntaxKind::Type
155                } else {
156                    SyntaxKind::Identifier
157                };
158                tokens.push(token(start..i, kind));
159            }
160            b'{' | b'}' | b'(' | b')' | b'[' | b']' | b',' | b';' | b':' | b'.' => {
161                tokens.push(token(i..i + 1, SyntaxKind::Punctuation));
162                i += 1;
163            }
164            b'+' | b'-' | b'*' | b'=' | b'!' | b'<' | b'>' | b'&' | b'|' | b'%' => {
165                let start = i;
166                i += 1;
167                while i < bytes.len() && matches!(bytes[i], b'=' | b'>' | b'&' | b'|') {
168                    i += 1;
169                }
170                tokens.push(token(start..i, SyntaxKind::Operator));
171            }
172            _ => i += 1,
173        }
174    }
175
176    mark_function_identifiers(text, &mut tokens);
177    mark_macro_identifiers(text, &mut tokens);
178    tokens
179}
180
181#[cfg(feature = "tree_sitter")]
182/// Combines Tree-sitter Rust tokens with lexical gap filling.
183///
184/// Returns `None` if parser setup/parsing fails or non-whitespace input yields no
185/// structural token. Otherwise ranges are clamped to UTF-8 boundaries, sorted,
186/// deduplicated, and made non-overlapping by semantic priority before and after
187/// lexical gap filling. Whitespace-only input returns `Some(Vec::new())`.
188///
189/// # Examples
190///
191/// ```
192/// use ailloli_ui_editor::code::{highlight_rust_tree_sitter_hybrid, SyntaxKind};
193/// let text = "fn main() { let n = 1; }";
194/// let tokens = highlight_rust_tree_sitter_hybrid(text).unwrap();
195/// assert!(tokens.iter().any(|t| t.kind == SyntaxKind::Function && &text[t.range.clone()] == "main"));
196/// assert_eq!(highlight_rust_tree_sitter_hybrid("   "), Some(Vec::new()));
197/// ```
198pub fn highlight_rust_tree_sitter_hybrid(text: &str) -> Option<Vec<SyntaxToken>> {
199    let mut parser = tree_sitter::Parser::new();
200    let language = tree_sitter_rust::LANGUAGE.into();
201    parser.set_language(&language).ok()?;
202    let tree = parser.parse(text, None)?;
203
204    let mut tokens = Vec::new();
205    collect_tree_sitter_tokens(tree.root_node(), &mut tokens);
206    if tokens.is_empty() && !text.trim().is_empty() {
207        return None;
208    }
209
210    let mut normalized = normalize_syntax_tokens(text, tokens);
211    gap_fill_lexical_tokens(text, &mut normalized);
212    Some(normalize_syntax_tokens(text, normalized))
213}
214
215#[cfg(feature = "tree_sitter")]
216/// Recursively collects structural and named-item tokens from a syntax tree.
217fn collect_tree_sitter_tokens(node: tree_sitter::Node<'_>, tokens: &mut Vec<SyntaxToken>) {
218    if let Some(kind) = tree_sitter_node_kind(node) {
219        tokens.push(token(node.start_byte()..node.end_byte(), kind));
220    }
221
222    if let Some((kind, name)) = tree_sitter_named_item(node) {
223        tokens.push(token(name.start_byte()..name.end_byte(), kind));
224    }
225
226    let mut cursor = node.walk();
227    for child in node.children(&mut cursor) {
228        collect_tree_sitter_tokens(child, tokens);
229    }
230}
231
232#[cfg(feature = "tree_sitter")]
233/// Maps a structural Tree-sitter Rust node to a syntax category.
234fn tree_sitter_node_kind(node: tree_sitter::Node<'_>) -> Option<SyntaxKind> {
235    match node.kind() {
236        "line_comment" | "block_comment" => Some(SyntaxKind::Comment),
237        "string_literal" | "raw_string_literal" | "char_literal" => Some(SyntaxKind::String),
238        "integer_literal" | "float_literal" => Some(SyntaxKind::Number),
239        "primitive_type" | "type_identifier" => Some(SyntaxKind::Type),
240        "visibility_modifier" | "where_clause" | "for_lifetimes" | "attribute_item" => {
241            Some(SyntaxKind::Keyword)
242        }
243        kind if RUST_KEYWORDS.contains(&kind) => Some(SyntaxKind::Keyword),
244        _ => None,
245    }
246}
247
248#[cfg(feature = "tree_sitter")]
249/// Extracts the named identifier node for declarations with semantic names.
250fn tree_sitter_named_item(
251    node: tree_sitter::Node<'_>,
252) -> Option<(SyntaxKind, tree_sitter::Node<'_>)> {
253    let kind = match node.kind() {
254        "function_item" => SyntaxKind::Function,
255        "struct_item" | "enum_item" | "trait_item" | "type_item" => SyntaxKind::Type,
256        _ => return None,
257    };
258    node.child_by_field_name("name").map(|name| (kind, name))
259}
260
261#[cfg(feature = "tree_sitter")]
262/// Adds lexical tokens not superseded by stronger structural tokens.
263fn gap_fill_lexical_tokens(text: &str, tokens: &mut Vec<SyntaxToken>) {
264    let existing = tokens.clone();
265    tokens.extend(
266        highlight_rust_lexical(text)
267            .into_iter()
268            .filter(|candidate| lexical_gap_fill_allowed(candidate, &existing)),
269    );
270}
271
272#[cfg(feature = "tree_sitter")]
273/// Tests whether a lexical token can coexist with current structural coverage.
274fn lexical_gap_fill_allowed(candidate: &SyntaxToken, existing: &[SyntaxToken]) -> bool {
275    let has_overlap = existing
276        .iter()
277        .any(|token| ranges_overlap(&candidate.range, &token.range));
278    if !has_overlap {
279        return true;
280    }
281    matches!(
282        candidate.kind,
283        SyntaxKind::Operator | SyntaxKind::Punctuation | SyntaxKind::Identifier
284    ) && !existing.iter().any(|token| {
285        ranges_overlap(&candidate.range, &token.range)
286            && syntax_kind_priority(token.kind) < syntax_kind_priority(candidate.kind)
287    })
288}
289
290#[cfg(feature = "tree_sitter")]
291/// Clamps, sorts, deduplicates, and resolves overlapping syntax tokens.
292fn normalize_syntax_tokens(text: &str, tokens: Vec<SyntaxToken>) -> Vec<SyntaxToken> {
293    let mut indexed = tokens
294        .into_iter()
295        .enumerate()
296        .filter_map(|(idx, token)| normalize_token(text, token).map(|token| (idx, token)))
297        .collect::<Vec<_>>();
298    indexed.sort_by(|(a_idx, a), (b_idx, b)| {
299        a.range
300            .start
301            .cmp(&b.range.start)
302            .then(a.range.end.cmp(&b.range.end))
303            .then(syntax_kind_priority(a.kind).cmp(&syntax_kind_priority(b.kind)))
304            .then(a_idx.cmp(b_idx))
305    });
306
307    let mut out: Vec<SyntaxToken> = Vec::new();
308    for (_, token) in indexed {
309        if out
310            .last()
311            .is_some_and(|prev| prev.range == token.range && prev.kind == token.kind)
312        {
313            continue;
314        }
315        if let Some(pos) = out
316            .iter()
317            .position(|prev| ranges_overlap(&prev.range, &token.range))
318        {
319            let previous_priority = syntax_kind_priority(out[pos].kind);
320            let current_priority = syntax_kind_priority(token.kind);
321            if current_priority < previous_priority {
322                out.remove(pos);
323                out.push(token);
324                out.sort_by(|a, b| {
325                    a.range
326                        .start
327                        .cmp(&b.range.start)
328                        .then(a.range.end.cmp(&b.range.end))
329                });
330            }
331            continue;
332        }
333        out.push(token);
334    }
335    out
336}
337
338#[cfg(feature = "tree_sitter")]
339/// Clamps one token to a non-empty valid UTF-8 source range.
340fn normalize_token(text: &str, token: SyntaxToken) -> Option<SyntaxToken> {
341    let start = clamp_to_char_boundary(text, token.range.start.min(text.len()));
342    let end = clamp_to_char_boundary(text, token.range.end.min(text.len()));
343    (end > start).then_some(SyntaxToken {
344        range: start..end,
345        kind: token.kind,
346    })
347}
348
349#[cfg(feature = "tree_sitter")]
350/// Moves an offset backward to the nearest UTF-8 boundary.
351fn clamp_to_char_boundary(text: &str, mut at: usize) -> usize {
352    while at > 0 && !text.is_char_boundary(at) {
353        at -= 1;
354    }
355    at
356}
357
358#[cfg(feature = "tree_sitter")]
359/// Tests intersection of two half-open byte ranges.
360fn ranges_overlap(a: &Range<usize>, b: &Range<usize>) -> bool {
361    a.start < b.end && b.start < a.end
362}
363
364#[cfg(feature = "tree_sitter")]
365/// Returns the lower-is-stronger overlap priority for a syntax category.
366fn syntax_kind_priority(kind: SyntaxKind) -> u8 {
367    match kind {
368        SyntaxKind::Comment => 0,
369        SyntaxKind::String => 1,
370        SyntaxKind::Number => 2,
371        SyntaxKind::Keyword => 3,
372        SyntaxKind::Function => 4,
373        SyntaxKind::Type => 5,
374        SyntaxKind::Identifier => 6,
375        SyntaxKind::Operator => 7,
376        SyntaxKind::Punctuation => 8,
377    }
378}
379
380/// Reclassifies the next identifier token after each `fn` keyword.
381fn mark_function_identifiers(text: &str, tokens: &mut [SyntaxToken]) {
382    for idx in 0..tokens.len().saturating_sub(1) {
383        if tokens[idx].kind == SyntaxKind::Keyword && &text[tokens[idx].range.clone()] == "fn" {
384            if let Some(next) = tokens[idx + 1..]
385                .iter_mut()
386                .find(|token| token.kind == SyntaxKind::Identifier)
387            {
388                next.kind = SyntaxKind::Function;
389            }
390        }
391    }
392}
393
394/// Reclassifies identifiers followed by optional whitespace and `!` as macros.
395fn mark_macro_identifiers(text: &str, tokens: &mut [SyntaxToken]) {
396    for token in tokens.iter_mut() {
397        if token.kind != SyntaxKind::Identifier {
398            continue;
399        }
400        let mut at = token.range.end;
401        while at < text.len() && text.as_bytes()[at].is_ascii_whitespace() {
402            at += 1;
403        }
404        if text.as_bytes().get(at) == Some(&b'!') {
405            token.kind = SyntaxKind::Function;
406        }
407    }
408}
409
410/// Returns the end of a Rust-like lifetime, rejecting character literals.
411fn lifetime_end(bytes: &[u8], start: usize) -> Option<usize> {
412    let first = *bytes.get(start + 1)?;
413    if !is_ident_start(first) {
414        return None;
415    }
416    let mut end = start + 2;
417    while end < bytes.len() && is_ident_continue(bytes[end]) {
418        end += 1;
419    }
420    if bytes.get(end) == Some(&b'\'') {
421        return None;
422    }
423    Some(end)
424}
425
426/// Returns the end of a Rust raw string, using EOF for an unterminated string.
427fn raw_string_end(text: &str, start: usize) -> Option<usize> {
428    let bytes = text.as_bytes();
429    if bytes.get(start) != Some(&b'r') {
430        return None;
431    }
432    let mut hashes = 0usize;
433    let mut quote = start + 1;
434    while bytes.get(quote) == Some(&b'#') {
435        hashes += 1;
436        quote += 1;
437    }
438    if bytes.get(quote) != Some(&b'"') {
439        return None;
440    }
441    let mut i = quote + 1;
442    while i < bytes.len() {
443        if bytes[i] == b'"' {
444            let mut matched = true;
445            for offset in 0..hashes {
446                if bytes.get(i + 1 + offset) != Some(&b'#') {
447                    matched = false;
448                    break;
449                }
450            }
451            if matched {
452                return Some((i + 1 + hashes).min(text.len()));
453            }
454        }
455        i += 1;
456    }
457    Some(text.len())
458}
459
460/// Constructs a syntax token without normalizing its range.
461fn token(range: Range<usize>, kind: SyntaxKind) -> SyntaxToken {
462    SyntaxToken { range, kind }
463}
464
465/// Tests the first byte of an ASCII Rust-like identifier.
466fn is_ident_start(ch: u8) -> bool {
467    ch == b'_' || ch.is_ascii_alphabetic()
468}
469
470/// Tests a continuation byte of an ASCII Rust-like identifier.
471fn is_ident_continue(ch: u8) -> bool {
472    is_ident_start(ch) || ch.is_ascii_digit()
473}