Skip to main content

oxframe_text/
code.rs

1//! A small, dependency-free syntax highlighter for code in animations.
2//! It favors robustness over completeness: unknown input degrades to plain
3//! text, never fails.
4
5use std::ops::Range;
6
7/// Language for syntax highlighting.
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
9pub enum Lang {
10    /// Rust.
11    Rust,
12    /// JavaScript and TypeScript.
13    Js,
14    /// Python.
15    Python,
16    /// JSON.
17    Json,
18    /// TOML.
19    Toml,
20    /// Shell scripts.
21    Shell,
22    /// No highlighting.
23    Plain,
24}
25
26impl Lang {
27    /// From a name or file extension (`"rs"`, `"typescript"`…); unknown names are `Plain`.
28    pub fn parse(name: &str) -> Self {
29        match name.to_ascii_lowercase().as_str() {
30            "rust" | "rs" => Self::Rust,
31            "js" | "javascript" | "ts" | "typescript" | "jsx" | "tsx" => Self::Js,
32            "py" | "python" => Self::Python,
33            "json" => Self::Json,
34            "toml" => Self::Toml,
35            "sh" | "bash" | "shell" | "zsh" => Self::Shell,
36            _ => Self::Plain,
37        }
38    }
39}
40
41/// Kind of a highlighted piece of code.
42#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
43pub enum Token {
44    /// Anything else.
45    Plain,
46    /// Keyword.
47    Keyword,
48    /// Type name.
49    Type,
50    /// Function name.
51    Function,
52    /// Macro name.
53    Macro,
54    /// String literal.
55    String,
56    /// Number literal.
57    Number,
58    /// Comment.
59    Comment,
60    /// Attribute or decorator.
61    Attribute,
62    /// Punctuation or operator.
63    Punct,
64}
65
66const RUST: &[&str] = &[
67    "as", "async", "await", "break", "const", "continue", "crate", "dyn", "else", "enum", "extern",
68    "false", "fn", "for", "if", "impl", "in", "let", "loop", "match", "mod", "move", "mut", "pub",
69    "ref", "return", "self", "Self", "static", "struct", "super", "trait", "true", "type",
70    "unsafe", "use", "where", "while",
71];
72const JS: &[&str] = &[
73    "as",
74    "async",
75    "await",
76    "break",
77    "case",
78    "catch",
79    "class",
80    "const",
81    "continue",
82    "default",
83    "delete",
84    "do",
85    "else",
86    "export",
87    "extends",
88    "false",
89    "finally",
90    "for",
91    "from",
92    "function",
93    "if",
94    "import",
95    "in",
96    "instanceof",
97    "interface",
98    "let",
99    "new",
100    "null",
101    "of",
102    "return",
103    "static",
104    "super",
105    "switch",
106    "this",
107    "throw",
108    "true",
109    "try",
110    "type",
111    "typeof",
112    "undefined",
113    "var",
114    "void",
115    "while",
116    "yield",
117];
118const PY: &[&str] = &[
119    "and", "as", "assert", "async", "await", "break", "class", "continue", "def", "del", "elif",
120    "else", "except", "False", "finally", "for", "from", "global", "if", "import", "in", "is",
121    "lambda", "None", "nonlocal", "not", "or", "pass", "raise", "return", "self", "True", "try",
122    "while", "with", "yield",
123];
124const SH: &[&str] = &[
125    "if", "then", "else", "elif", "fi", "for", "in", "do", "done", "while", "case", "esac",
126    "function", "return", "export", "local", "echo", "cd",
127];
128
129/// Splits `code` into highlighted ranges covering all of it.
130pub fn highlight(code: &str, lang: Lang) -> Vec<(Range<usize>, Token)> {
131    let b = code.as_bytes();
132    let mut out: Vec<(Range<usize>, Token)> = Vec::new();
133    let push = |out: &mut Vec<(Range<usize>, Token)>, r: Range<usize>, t: Token| {
134        if r.is_empty() {
135            return;
136        }
137        match out.last_mut() {
138            Some((last, lt)) if *lt == t && last.end == r.start => last.end = r.end,
139            _ => out.push((r, t)),
140        }
141    };
142    if lang == Lang::Plain {
143        push(&mut out, 0..code.len(), Token::Plain);
144        return out;
145    }
146    let keywords: &[&str] = match lang {
147        Lang::Rust => RUST,
148        Lang::Js => JS,
149        Lang::Python => PY,
150        Lang::Shell => SH,
151        Lang::Json => &["true", "false", "null"],
152        Lang::Toml => &["true", "false"],
153        Lang::Plain => &[],
154    };
155    let line_comment: &[&str] = match lang {
156        Lang::Rust | Lang::Js => &["//"],
157        Lang::Python | Lang::Shell | Lang::Toml => &["#"],
158        _ => &[],
159    };
160    let block_comment = matches!(lang, Lang::Rust | Lang::Js);
161    let mut i = 0;
162    while i < b.len() {
163        let rest = &code[i..];
164        let c = b[i];
165        // Attributes: #[...] in Rust, @decorator in Python/TS.
166        if lang == Lang::Rust && rest.starts_with("#[")
167            || lang == Lang::Rust && rest.starts_with("#![")
168        {
169            let end = rest.find(']').map_or(code.len(), |e| i + e + 1);
170            push(&mut out, i..end, Token::Attribute);
171            i = end;
172            continue;
173        }
174        if matches!(lang, Lang::Python | Lang::Js) && c == b'@' {
175            let end = i
176                + 1
177                + rest[1..]
178                    .find(|ch: char| !(ch.is_alphanumeric() || ch == '_' || ch == '.'))
179                    .unwrap_or(rest.len() - 1);
180            push(&mut out, i..end, Token::Attribute);
181            i = end;
182            continue;
183        }
184        if let Some(lc) = line_comment.iter().find(|p| rest.starts_with(**p)) {
185            let _ = lc;
186            let end = rest.find('\n').map_or(code.len(), |e| i + e);
187            push(&mut out, i..end, Token::Comment);
188            i = end;
189            continue;
190        }
191        if block_comment && rest.starts_with("/*") {
192            let end = rest[2..].find("*/").map_or(code.len(), |e| i + 2 + e + 2);
193            push(&mut out, i..end, Token::Comment);
194            i = end;
195            continue;
196        }
197        // TOML table headers.
198        if lang == Lang::Toml && c == b'[' && (i == 0 || b[i - 1] == b'\n') {
199            let end = rest.find('\n').map_or(code.len(), |e| i + e);
200            push(&mut out, i..end, Token::Type);
201            i = end;
202            continue;
203        }
204        if c == b'"' || c == b'\'' || (c == b'`' && lang == Lang::Js) {
205            // Rust lifetimes ('a) are not strings.
206            if c == b'\'' && lang == Lang::Rust {
207                let mut j = i + 1;
208                while j < b.len() && (b[j].is_ascii_alphanumeric() || b[j] == b'_') {
209                    j += 1;
210                }
211                if j < b.len() && b[j] != b'\'' && j > i + 1 {
212                    push(&mut out, i..j, Token::Type);
213                    i = j;
214                    continue;
215                }
216            }
217            let mut j = i + 1;
218            while j < b.len() && b[j] != c {
219                if b[j] == b'\\' {
220                    j += 1;
221                }
222                if b[j.min(b.len() - 1)] == b'\n' && c != b'`' && c != b'"' {
223                    break;
224                }
225                j += 1;
226            }
227            let end = (j + 1).min(b.len());
228            push(&mut out, i..end, Token::String);
229            i = end;
230            continue;
231        }
232        if c.is_ascii_digit() {
233            let mut j = i;
234            while j < b.len() && (b[j].is_ascii_alphanumeric() || b[j] == b'_' || b[j] == b'.') {
235                // `1..2` is a range, not a number.
236                if b[j] == b'.' && j + 1 < b.len() && b[j + 1] == b'.' {
237                    break;
238                }
239                j += 1;
240            }
241            push(&mut out, i..j, Token::Number);
242            i = j;
243            continue;
244        }
245        if c.is_ascii_alphabetic() || c == b'_' || c >= 0x80 {
246            let mut j = i;
247            while j < b.len() && (b[j].is_ascii_alphanumeric() || b[j] == b'_' || b[j] >= 0x80) {
248                j += 1;
249            }
250            let word = &code[i..j];
251            let next = code[j..].chars().find(|ch| !ch.is_whitespace());
252            let t = if keywords.contains(&word) {
253                Token::Keyword
254            } else if lang == Lang::Rust && code[j..].starts_with('!') {
255                Token::Macro
256            } else if lang == Lang::Json && next == Some(':') {
257                Token::Attribute
258            } else if next == Some('(') || (lang == Lang::Rust && code[j..].starts_with("::<")) {
259                Token::Function
260            } else if word.chars().next().is_some_and(char::is_uppercase) {
261                Token::Type
262            } else if lang == Lang::Toml && next == Some('=') {
263                Token::Attribute
264            } else {
265                Token::Plain
266            };
267            let end = if t == Token::Macro { j + 1 } else { j };
268            push(&mut out, i..end, t);
269            i = end;
270            continue;
271        }
272        let ch_len = rest.chars().next().map_or(1, char::len_utf8);
273        let t = if c.is_ascii_punctuation() {
274            Token::Punct
275        } else {
276            Token::Plain
277        };
278        push(&mut out, i..i + ch_len, t);
279        i += ch_len;
280    }
281    out
282}
283
284#[cfg(test)]
285mod tests {
286    use super::*;
287
288    fn tokens(code: &str, lang: Lang) -> Vec<(&str, Token)> {
289        highlight(code, lang)
290            .into_iter()
291            .filter(|(_, t)| !matches!(t, Token::Plain | Token::Punct))
292            .map(|(r, t)| (&code[r], t))
293            .collect()
294    }
295
296    #[test]
297    fn rust() {
298        let code =
299            "#[oxframe::scene]\nasync fn intro(v: View) { println!(\"hi\"); let x = 1.5; } // done";
300        let t = tokens(code, Lang::Rust);
301        assert!(t.contains(&("#[oxframe::scene]", Token::Attribute)));
302        assert!(t.contains(&("async", Token::Keyword)));
303        assert!(t.contains(&("intro", Token::Function)));
304        assert!(t.contains(&("View", Token::Type)));
305        assert!(t.contains(&("println!", Token::Macro)));
306        assert!(t.contains(&("\"hi\"", Token::String)));
307        assert!(t.contains(&("1.5", Token::Number)));
308        assert!(t.contains(&("// done", Token::Comment)));
309    }
310
311    #[test]
312    fn covers_everything_and_never_panics() {
313        for (code, lang) in [
314            ("const a = `x${y}`; // c", Lang::Js),
315            ("def f(x):\n    return 'a' # c", Lang::Python),
316            ("{\"a\": [1, true]}", Lang::Json),
317            ("[dev]\nport = 3000", Lang::Toml),
318            ("echo \"$HOME\" # c", Lang::Shell),
319            ("'unterminated", Lang::Rust),
320            ("\"open", Lang::Js),
321            ("Привет мир", Lang::Rust),
322            ("for i in 0..10 {}", Lang::Rust),
323        ] {
324            let h = highlight(code, lang);
325            let joined: String = h.iter().map(|(r, _)| &code[r.clone()]).collect();
326            assert_eq!(joined, code, "{lang:?}");
327        }
328        assert!(tokens("for i in 0..10 {}", Lang::Rust).contains(&("0", Token::Number)));
329    }
330}