1use std::ops::Range;
6
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
9pub enum Lang {
10 Rust,
12 Js,
14 Python,
16 Json,
18 Toml,
20 Shell,
22 Plain,
24}
25
26impl Lang {
27 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#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
43pub enum Token {
44 Plain,
46 Keyword,
48 Type,
50 Function,
52 Macro,
54 String,
56 Number,
58 Comment,
60 Attribute,
62 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
129pub 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 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 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 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 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}