Skip to main content

run/
highlight.rs

1use once_cell::sync::Lazy;
2use syntect::easy::HighlightLines;
3use syntect::highlighting::ThemeSet;
4use syntect::parsing::{SyntaxReference, SyntaxSet};
5use syntect::util::{LinesWithEndings, as_24_bit_terminal_escaped};
6
7static SYNTAX_SET: Lazy<SyntaxSet> = Lazy::new(SyntaxSet::load_defaults_newlines);
8
9static THEME_SET: Lazy<ThemeSet> = Lazy::new(ThemeSet::load_defaults);
10
11fn supports_color() -> bool {
12    std::env::var("NO_COLOR").is_err()
13        && (std::env::var("TERM").unwrap_or_default().contains("color")
14            || std::env::var("COLORTERM").is_ok())
15}
16
17fn language_to_syntax_name(language_id: &str) -> &str {
18    match language_id.to_ascii_lowercase().as_str() {
19        "python" | "py" | "python3" | "py3" => "Python",
20        "javascript" | "js" | "node" | "nodejs" => "JavaScript",
21        "typescript" | "ts" => "JavaScript",
22        "rust" | "rs" => "Rust",
23        "go" | "golang" => "Go",
24        "c" => "C",
25        "cpp" | "c++" | "cxx" => "C++",
26        "java" => "Java",
27        "csharp" | "cs" | "c#" => "C#",
28        "ruby" | "rb" => "Ruby",
29        "php" => "PHP",
30        "bash" | "sh" | "shell" | "zsh" => "Bourne Again Shell (bash)",
31        "lua" => "Lua",
32        "perl" | "pl" => "Perl",
33        "swift" => "Swift",
34        "kotlin" | "kt" => "Kotlin",
35        "r" | "rscript" => "R",
36        "haskell" | "hs" => "Haskell",
37        "julia" | "jl" => "Plain Text",
38        "elixir" | "ex" | "exs" => "Plain Text",
39        "dart" => "Dart",
40        "groovy" | "grv" => "Groovy",
41        "crystal" | "cr" => "Crystal",
42        "zig" => "Zig",
43        "nim" => "Nim",
44        _ => "Plain Text",
45    }
46}
47
48fn get_syntax_for_language(language_id: &str) -> &'static SyntaxReference {
49    let syntax_name = language_to_syntax_name(language_id);
50    SYNTAX_SET
51        .find_syntax_by_name(syntax_name)
52        .unwrap_or_else(|| SYNTAX_SET.find_syntax_plain_text())
53}
54
55pub fn highlight_code(code: &str, language_id: &str) -> String {
56    if !supports_color() {
57        return code.to_string();
58    }
59
60    let syntax = get_syntax_for_language(language_id);
61
62    let theme = &THEME_SET.themes["base16-ocean.dark"];
63
64    let mut highlighter = HighlightLines::new(syntax, theme);
65    let mut output = String::new();
66
67    for line in LinesWithEndings::from(code) {
68        let ranges = highlighter
69            .highlight_line(line, &SYNTAX_SET)
70            .unwrap_or_default();
71        let escaped = as_24_bit_terminal_escaped(&ranges[..], false);
72        output.push_str(&escaped);
73    }
74
75    if !output.is_empty() && !output.ends_with("\x1b[0m") {
76        output.push_str("\x1b[0m");
77    }
78
79    output
80}
81
82pub fn highlight_repl_input(code: &str, language_id: &str) -> String {
83    highlight_code(code, language_id)
84}
85
86pub fn highlight_output(code: &str, language_id: &str) -> String {
87    highlight_code(code, language_id)
88}
89
90pub fn has_syntax_support(language_id: &str) -> bool {
91    let syntax_name = language_to_syntax_name(language_id);
92    SYNTAX_SET.find_syntax_by_name(syntax_name).is_some()
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98
99    #[test]
100    fn test_language_mapping() {
101        assert_eq!(language_to_syntax_name("python"), "Python");
102        assert_eq!(language_to_syntax_name("rust"), "Rust");
103        assert_eq!(language_to_syntax_name("javascript"), "JavaScript");
104        assert_eq!(language_to_syntax_name("typescript"), "JavaScript");
105        assert_eq!(language_to_syntax_name("go"), "Go");
106        assert_eq!(language_to_syntax_name("java"), "Java");
107        assert_eq!(language_to_syntax_name("csharp"), "C#");
108        assert_eq!(language_to_syntax_name("cpp"), "C++");
109        assert_eq!(language_to_syntax_name("ruby"), "Ruby");
110        assert_eq!(language_to_syntax_name("php"), "PHP");
111    }
112
113    #[test]
114    fn test_all_language_aliases() {
115        assert_eq!(language_to_syntax_name("py"), "Python");
116        assert_eq!(language_to_syntax_name("py3"), "Python");
117        assert_eq!(language_to_syntax_name("python3"), "Python");
118
119        assert_eq!(language_to_syntax_name("js"), "JavaScript");
120        assert_eq!(language_to_syntax_name("node"), "JavaScript");
121        assert_eq!(language_to_syntax_name("nodejs"), "JavaScript");
122
123        assert_eq!(language_to_syntax_name("ts"), "JavaScript");
124
125        assert_eq!(language_to_syntax_name("rs"), "Rust");
126
127        assert_eq!(language_to_syntax_name("golang"), "Go");
128
129        assert_eq!(language_to_syntax_name("c++"), "C++");
130        assert_eq!(language_to_syntax_name("cxx"), "C++");
131
132        assert_eq!(language_to_syntax_name("cs"), "C#");
133        assert_eq!(language_to_syntax_name("c#"), "C#");
134    }
135
136    #[test]
137    fn test_all_languages_have_syntax() {
138        let supported = vec![
139            "python",
140            "javascript",
141            "rust",
142            "go",
143            "c",
144            "cpp",
145            "java",
146            "csharp",
147            "ruby",
148            "php",
149            "bash",
150            "lua",
151            "perl",
152            "r",
153            "haskell",
154        ];
155
156        for lang in supported {
157            assert!(
158                has_syntax_support(lang),
159                "Language {} should be supported",
160                lang
161            );
162        }
163
164        let fallback = vec![
165            "swift", "kotlin", "dart", "groovy", "crystal", "zig", "nim", "julia", "elixir",
166        ];
167        for lang in fallback {
168            let _ = highlight_code("test", lang);
169        }
170    }
171
172    #[test]
173    fn test_unknown_language_fallback() {
174        let syntax_name = language_to_syntax_name("unknownlang123");
175        assert_eq!(syntax_name, "Plain Text");
176    }
177
178    #[test]
179    fn test_syntax_available() {
180        assert!(has_syntax_support("python"));
181        assert!(has_syntax_support("rust"));
182        assert!(has_syntax_support("go"));
183        assert!(has_syntax_support("javascript"));
184        assert!(has_syntax_support("typescript"));
185    }
186
187    #[test]
188    fn test_highlight_basic() {
189        let code = "fn main() { println!(\"Hello\"); }";
190        let highlighted = highlight_code(code, "rust");
191        assert!(!highlighted.is_empty());
192        assert!(highlighted.len() >= code.len());
193    }
194
195    #[test]
196    fn test_highlight_python() {
197        let code = "def hello():\n    print('world')";
198        let highlighted = highlight_code(code, "python");
199        assert!(!highlighted.is_empty());
200    }
201
202    #[test]
203    fn test_highlight_javascript() {
204        let code = "function hello() { console.log('world'); }";
205        let highlighted = highlight_code(code, "javascript");
206        assert!(!highlighted.is_empty());
207    }
208
209    #[test]
210    fn test_highlight_go() {
211        let code = "package main\nfunc main() { fmt.Println(\"hello\") }";
212        let highlighted = highlight_code(code, "go");
213        assert!(!highlighted.is_empty());
214    }
215
216    #[test]
217    fn test_empty_code() {
218        let highlighted = highlight_code("", "python");
219        assert!(highlighted.is_empty() || highlighted == "\x1b[0m");
220    }
221
222    #[test]
223    fn test_whitespace_only() {
224        let code = "   \n  \t  ";
225        let highlighted = highlight_code(code, "rust");
226        assert!(!highlighted.is_empty());
227    }
228
229    #[test]
230    fn test_multiline_code() {
231        let code = "fn main() {\n    let x = 10;\n    println!(\"{}\", x);\n}";
232        let highlighted = highlight_code(code, "rust");
233        assert!(!highlighted.is_empty());
234        assert!(highlighted.contains('\n') || highlighted.contains("\\n"));
235    }
236
237    #[test]
238    fn test_color_reset_at_end() {
239        unsafe {
240            std::env::set_var("TERM", "xterm-256color");
241        }
242
243        let code = "x = 10";
244        let highlighted = highlight_code(code, "python");
245
246        assert!(
247            highlighted.ends_with("\x1b[0m") || !highlighted.contains("\x1b["),
248            "Highlighted code should end with color reset or be plain text"
249        );
250    }
251
252    #[test]
253    fn test_no_color_environment() {
254        unsafe {
255            std::env::set_var("NO_COLOR", "1");
256        }
257
258        let code = "fn main() {}";
259        let highlighted = highlight_code(code, "rust");
260
261        assert_eq!(highlighted, code);
262
263        unsafe {
264            std::env::remove_var("NO_COLOR");
265        }
266    }
267
268    #[test]
269    fn test_special_characters() {
270        let code = "print(\"Hello\\nWorld\\t!\")";
271        let highlighted = highlight_code(code, "python");
272        assert!(!highlighted.is_empty());
273    }
274
275    #[test]
276    fn test_string_literal() {
277        let code = "message = \"Hello World\"";
278        let highlighted = highlight_code(code, "python");
279        assert!(!highlighted.is_empty());
280    }
281
282    #[test]
283    fn test_repl_input_helper() {
284        let code = "x = 42";
285        let highlighted = highlight_repl_input(code, "python");
286        assert!(!highlighted.is_empty());
287    }
288
289    #[test]
290    fn test_output_helper() {
291        let code = "console.log('test')";
292        let highlighted = highlight_output(code, "javascript");
293        assert!(!highlighted.is_empty());
294    }
295
296    #[test]
297    fn test_bash_highlighting() {
298        let code = "echo \"Hello World\"";
299        let highlighted = highlight_code(code, "bash");
300        assert!(!highlighted.is_empty());
301    }
302
303    #[test]
304    fn test_sql_like_code_in_supported_language() {
305        let code = "query = \"SELECT * FROM users\"";
306        let highlighted = highlight_code(code, "python");
307        assert!(!highlighted.is_empty());
308    }
309
310    #[test]
311    fn test_comments_highlighting() {
312        let code = "// This is a comment\nlet x = 10;";
313        let highlighted = highlight_code(code, "javascript");
314        assert!(!highlighted.is_empty());
315    }
316
317    #[test]
318    fn test_string_highlighting() {
319        let code = "\"This is a string\"";
320        let highlighted = highlight_code(code, "python");
321        assert!(!highlighted.is_empty());
322    }
323
324    #[test]
325    fn test_number_highlighting() {
326        let code = "let num = 42;";
327        let highlighted = highlight_code(code, "javascript");
328        assert!(!highlighted.is_empty());
329    }
330
331    #[test]
332    fn test_very_long_code() {
333        let mut code = String::new();
334        for i in 0..1000 {
335            code.push_str(&format!("let var{} = {};\n", i, i));
336        }
337        let highlighted = highlight_code(&code, "javascript");
338        assert!(!highlighted.is_empty());
339    }
340
341    #[test]
342    fn test_case_insensitive_language_names() {
343        assert_eq!(language_to_syntax_name("PYTHON"), "Python");
344        assert_eq!(language_to_syntax_name("Python"), "Python");
345        assert_eq!(language_to_syntax_name("RuSt"), "Rust");
346        assert_eq!(language_to_syntax_name("JavaScript"), "JavaScript");
347    }
348
349    #[test]
350    fn test_syntax_reference_caching() {
351        let syntax1 = get_syntax_for_language("python");
352        let syntax2 = get_syntax_for_language("python");
353        assert_eq!(syntax1.name, syntax2.name);
354    }
355
356    #[test]
357    fn test_all_functional_languages() {
358        let langs = vec!["haskell", "elixir", "julia"];
359        for lang in langs {
360            let code = "main = print \"hello\"";
361            let highlighted = highlight_code(code, lang);
362            assert!(!highlighted.is_empty(), "Failed for {}", lang);
363        }
364    }
365
366    #[test]
367    fn test_systems_languages() {
368        let langs = vec!["c", "cpp", "rust", "zig", "nim"];
369        for lang in langs {
370            let code = "int main() { return 0; }";
371            let highlighted = highlight_code(code, lang);
372            assert!(!highlighted.is_empty(), "Failed for {}", lang);
373        }
374    }
375
376    #[test]
377    fn test_scripting_languages() {
378        let langs = vec!["python", "ruby", "perl", "lua", "php"];
379        for lang in langs {
380            let code = "print('hello')";
381            let highlighted = highlight_code(code, lang);
382            assert!(!highlighted.is_empty(), "Failed for {}", lang);
383        }
384    }
385
386    #[test]
387    fn test_jvm_languages() {
388        let langs = vec!["java", "kotlin", "groovy"];
389        for lang in langs {
390            let code = "public class Test { }";
391            let highlighted = highlight_code(code, lang);
392            assert!(!highlighted.is_empty(), "Failed for {}", lang);
393        }
394    }
395
396    #[test]
397    fn test_ansi_codes_present_when_colors_enabled() {
398        unsafe {
399            std::env::set_var("TERM", "xterm-256color");
400            std::env::remove_var("NO_COLOR");
401        }
402
403        let code = "fn main() {}";
404        let highlighted = highlight_code(code, "rust");
405
406        let has_ansi = highlighted.contains("\x1b[") || highlighted == code;
407        assert!(has_ansi, "Should contain ANSI codes or be plain text");
408    }
409
410    #[test]
411    fn test_typescript_mapping() {
412        assert_eq!(language_to_syntax_name("typescript"), "JavaScript");
413        let code = "const x: number = 42;";
414        let highlighted = highlight_code(code, "typescript");
415        assert!(!highlighted.is_empty());
416    }
417
418    #[test]
419    fn test_swift_language() {
420        let code = "func main() { print(\"Hello\") }";
421        let highlighted = highlight_code(code, "swift");
422        assert!(!highlighted.is_empty());
423    }
424
425    #[test]
426    fn test_kotlin_language() {
427        let code = "fun main() { println(\"Hello\") }";
428        let highlighted = highlight_code(code, "kotlin");
429        assert!(!highlighted.is_empty());
430    }
431
432    #[test]
433    fn test_dart_language() {
434        let code = "void main() { print('Hello'); }";
435        let highlighted = highlight_code(code, "dart");
436        assert!(!highlighted.is_empty());
437    }
438
439    #[test]
440    fn test_r_language() {
441        let code = "x <- 10\nprint(x)";
442        let highlighted = highlight_code(code, "r");
443        assert!(!highlighted.is_empty());
444    }
445
446    #[test]
447    fn test_crystal_language() {
448        let code = "puts \"Hello\"";
449        let highlighted = highlight_code(code, "crystal");
450        assert!(!highlighted.is_empty());
451    }
452
453    #[test]
454    fn test_zig_language() {
455        let code = "pub fn main() void {}";
456        let highlighted = highlight_code(code, "zig");
457        assert!(!highlighted.is_empty());
458    }
459
460    #[test]
461    fn test_nim_language() {
462        let code = "echo \"Hello\"";
463        let highlighted = highlight_code(code, "nim");
464        assert!(!highlighted.is_empty());
465    }
466
467    #[test]
468    fn test_only_comments() {
469        let code = "// Just a comment\n// Another comment";
470        let highlighted = highlight_code(code, "rust");
471        assert!(!highlighted.is_empty());
472    }
473
474    #[test]
475    fn test_code_with_syntax_errors() {
476        let code = "fn main( { println! }";
477        let highlighted = highlight_code(code, "rust");
478        assert!(!highlighted.is_empty());
479    }
480
481    #[test]
482    fn test_mixed_quotes() {
483        let code = r#"s1 = 'single'; s2 = "double""#;
484        let highlighted = highlight_code(code, "python");
485        assert!(!highlighted.is_empty());
486    }
487
488    #[test]
489    fn test_nested_structures() {
490        let code = "let arr = [[1, 2], [3, 4], [5, 6]];";
491        let highlighted = highlight_code(code, "javascript");
492        assert!(!highlighted.is_empty());
493    }
494
495    #[test]
496    fn test_regex_patterns() {
497        let code = r"pattern = /[a-z]+/g";
498        let highlighted = highlight_code(code, "javascript");
499        assert!(!highlighted.is_empty());
500    }
501
502    #[test]
503    fn test_multiple_statements_one_line() {
504        let code = "x = 1; y = 2; z = 3;";
505        let highlighted = highlight_code(code, "javascript");
506        assert!(!highlighted.is_empty());
507    }
508
509    #[test]
510    fn test_color_reset_prevents_bleeding() {
511        unsafe {
512            std::env::set_var("TERM", "xterm-256color");
513            std::env::remove_var("NO_COLOR");
514        }
515
516        let code1 = "x = 20";
517        let highlighted1 = highlight_code(code1, "javascript");
518
519        if highlighted1.contains("\x1b[") {
520            assert!(
521                highlighted1.ends_with("\x1b[0m"),
522                "Highlighted code must end with reset code to prevent color bleeding"
523            );
524        }
525    }
526
527    #[test]
528    fn test_all_25_languages_work() {
529        let languages = vec![
530            "python",
531            "javascript",
532            "typescript",
533            "rust",
534            "go",
535            "c",
536            "cpp",
537            "java",
538            "csharp",
539            "ruby",
540            "php",
541            "bash",
542            "lua",
543            "perl",
544            "swift",
545            "kotlin",
546            "r",
547            "dart",
548            "haskell",
549            "julia",
550            "elixir",
551            "groovy",
552            "crystal",
553            "zig",
554            "nim",
555        ];
556
557        for lang in languages {
558            let code = "test code";
559            let highlighted = highlight_code(code, lang);
560            assert!(
561                !highlighted.is_empty(),
562                "Language {} failed to highlight",
563                lang
564            );
565        }
566    }
567
568    #[test]
569    fn test_bash_aliases() {
570        let aliases = vec!["bash", "sh", "shell", "zsh"];
571        for alias in aliases {
572            assert_eq!(
573                language_to_syntax_name(alias),
574                "Bourne Again Shell (bash)",
575                "Bash alias {} failed",
576                alias
577            );
578        }
579    }
580
581    #[test]
582    fn test_empty_lines_in_code() {
583        let code = "fn main() {\n\n\n    println!(\"test\");\n\n}";
584        let highlighted = highlight_code(code, "rust");
585        assert!(!highlighted.is_empty());
586
587        assert!(highlighted.matches('\n').count() >= 4);
588    }
589
590    #[test]
591    fn test_tabs_and_spaces() {
592        let code = "\tfn main() {\n\t\tprintln!(\"test\");\n\t}";
593        let highlighted = highlight_code(code, "rust");
594        assert!(!highlighted.is_empty());
595    }
596
597    #[test]
598    fn test_colorterm_environment() {
599        unsafe {
600            std::env::set_var("COLORTERM", "truecolor");
601            std::env::remove_var("NO_COLOR");
602        }
603
604        assert!(supports_color());
605
606        unsafe {
607            std::env::remove_var("COLORTERM");
608        }
609    }
610
611    #[test]
612    fn test_repl_helpers_consistency() {
613        let code = "test = 42";
614        let h1 = highlight_code(code, "python");
615        let h2 = highlight_repl_input(code, "python");
616        let h3 = highlight_output(code, "python");
617
618        assert_eq!(h1, h2);
619        assert_eq!(h2, h3);
620    }
621}