1use std::{
2 cell::RefCell,
3 collections::HashMap,
4 fmt::Write as FmtWrite,
5 io::{ErrorKind, Write as IoWrite},
6 path::Path,
7 sync::LazyLock,
8};
9
10use comrak::adapters::SyntaxHighlighterAdapter;
11use tracing::{debug, error};
12use tree_sitter_grammar_repository::{Grammar, Language};
13use tree_sitter_highlight::{HighlightConfiguration, HighlightEvent, Highlighter};
14
15thread_local! {
16 static HIGHLIGHTER: RefCell<Highlighter> = RefCell::new(Highlighter::new());
17}
18
19macro_rules! count {
20 () => (0);
21 ($e:expr) => (1);
22 ($e:expr, $($rest:expr),*) => (1 + count!($($rest),*));
23}
24
25macro_rules! define_classes {
26 ($($name:literal => $class:literal),*,) => {
27 static HIGHLIGHT_NAMES: [&str; count!($($name),*)] = [
28 $($name),*
29 ];
30
31 static HIGHLIGHT_CLASSES: [&str; count!($($name),*)] = [
32 $($class),*
33 ];
34 };
35}
36
37define_classes! {
38"attribute" => "attribute",
39"boolean" => "boolean",
40"carriage-return" => "carriage-return",
41"comment" => "comment",
42"comment.documentation" => "comment documentation",
43"constant" => "constant",
44"constant.builtin" => "constant builtin",
45"constructor" => "constructor",
46"constructor.builtin" => "constructor builtin",
47"embedded" => "embedded",
48"error" => "error",
49"escape" => "escape",
50"function" => "function",
51"function.builtin" => "function builtin",
52"keyword" => "keyword",
53"markup" => "markup",
54"markup.bold" => "markup bold",
55"markup.heading" => "markup heading",
56"markup.italic" => "markup italic",
57"markup.link" => "markup link",
58"markup.link.url" => "markup link url",
59"markup.list" => "markup list",
60"markup.list.checked" => "markup list checked",
61"markup.list.numbered" => "markup list numbered",
62"markup.list.unchecked" => "markup list unchecked",
63"markup.list.unnumbered" => "markup list unnumbered",
64"markup.quote" => "markup quote",
65"markup.raw" => "markup raw",
66"markup.raw.block" => "markup raw block",
67"markup.raw.inline" => "markup raw inline",
68"markup.strikethrough" => "markup strikethrough",
69"module" => "module",
70"number" => "number",
71"operator" => "operator",
72"property" => "property",
73"property.builtin" => "property builtin",
74"punctuation" => "punctuation",
75"punctuation.bracket" => "punctuation bracket",
76"punctuation.delimiter" => "punctuation delimiter",
77"punctuation.special" => "punctuation special",
78"string" => "string",
79"string.escape" => "string escape",
80"string.regexp" => "string regexp",
81"string.special" => "string special",
82"string.special.symbol" => "string special symbol",
83"tag" => "tag",
84"type" => "type",
85"type.builtin" => "type builtin",
86"variable" => "variable",
87"variable.builtin" => "variable builtin",
88"variable.member" => "variable member",
89"variable.parameter" => "variable parameter",}
90
91pub fn prime_highlighters() {
92 let _res = HIGHLIGHTER_CONFIGS.len();
93}
94
95static HIGHLIGHTER_CONFIGS: LazyLock<Vec<HighlightConfiguration>> = LazyLock::new(|| {
96 Grammar::VARIANTS
97 .iter()
98 .copied()
99 .map(Grammar::highlight_configuration_params)
100 .map(|v| {
101 let mut configuration = HighlightConfiguration::new(
102 v.language.into(),
103 v.name,
104 v.highlights_query,
105 v.injection_query,
106 v.locals_query,
107 )
108 .unwrap_or_else(|e| panic!("bad query for {}: {e}", v.name));
109 configuration.configure(&HIGHLIGHT_NAMES);
110 configuration
111 })
112 .collect()
113});
114
115pub fn fetch_highlighter_config(file: &Path) -> Option<&'static HighlightConfiguration> {
116 Language::from_file_name(file)
117 .map(Language::grammar)
118 .map(Grammar::idx)
119 .map(|idx| &HIGHLIGHTER_CONFIGS[idx])
120}
121
122pub fn fetch_highlighter_config_by_token(token: &str) -> Option<&'static HighlightConfiguration> {
123 Language::from_injection(token)
124 .map(Language::grammar)
125 .map(Grammar::idx)
126 .map(|idx| &HIGHLIGHTER_CONFIGS[idx])
127}
128
129pub struct ComrakHighlightAdapter;
130
131impl SyntaxHighlighterAdapter for ComrakHighlightAdapter {
132 fn write_highlighted(
133 &self,
134 output: &mut dyn IoWrite,
135 lang: Option<&str>,
136 code: &str,
137 ) -> std::io::Result<()> {
138 let out = format_file(code, FileIdentifier::Token(lang.unwrap_or_default()))
139 .map_err(|e| std::io::Error::new(ErrorKind::Other, e))?;
140 output.write_all(out.as_bytes())
141 }
142
143 fn write_pre_tag(
144 &self,
145 output: &mut dyn IoWrite,
146 _attributes: HashMap<String, String>,
147 ) -> std::io::Result<()> {
148 write!(output, "<pre>")
149 }
150
151 fn write_code_tag(
152 &self,
153 _output: &mut dyn IoWrite,
154 _attributes: HashMap<String, String>,
155 ) -> std::io::Result<()> {
156 Ok(())
157 }
158}
159
160#[derive(Copy, Clone, Debug)]
161pub enum FileIdentifier<'a> {
162 Path(&'a Path),
163 Token(&'a str),
164}
165
166pub fn format_file(content: &str, identifier: FileIdentifier<'_>) -> anyhow::Result<String> {
167 let mut out = String::new();
168 format_file_inner(&mut out, content, identifier, true)?;
169 Ok(out)
170}
171
172pub fn format_file_inner(
173 out: &mut String,
174 content: &str,
175 identifier: FileIdentifier<'_>,
176 code_tag: bool,
177) -> anyhow::Result<()> {
178 let config = match identifier {
179 FileIdentifier::Path(v) => fetch_highlighter_config(v),
180 FileIdentifier::Token(v) => fetch_highlighter_config_by_token(v),
181 };
182
183 let line_prefix = if code_tag { "<code>" } else { "" };
184 let line_suffix = if code_tag { "</code>\n" } else { "\n" };
185
186 let Some(config) = config else {
187 for line in content.lines() {
188 out.push_str(line_prefix);
189 v_htmlescape::b_escape(line.as_bytes(), out);
190 out.push_str(line_suffix);
191 }
192
193 return Ok(());
194 };
195
196 HIGHLIGHTER.with_borrow_mut(|highlighter| {
197 highlighter.parser().reset();
198
199 let spans = highlighter.highlight(config, content.as_bytes(), None, |injection| {
200 debug!(injection, "Highlighter switch requested");
201 fetch_highlighter_config_by_token(injection)
202 });
203
204 let mut spans = match spans {
205 Ok(v) => v,
206 Err(error) => {
207 error!(
208 ?error,
209 "Failed to run highlighter, falling back to plaintext"
210 );
211
212 for line in content.lines() {
213 out.push_str(line_prefix);
214 v_htmlescape::b_escape(line.as_bytes(), out);
215 out.push_str(line_suffix);
216 }
217
218 return Ok(());
219 }
220 };
221
222 let mut tag_open = true;
223 out.push_str(line_prefix);
224
225 while let Some(span) = spans.next().transpose()? {
226 if !tag_open {
227 out.push_str(line_prefix);
228 tag_open = true;
229 }
230
231 match span {
232 HighlightEvent::Source { start, end } => {
233 let content = &content[start..end];
234
235 for (i, line) in content.lines().enumerate() {
236 if i != 0 {
237 out.push_str(line_suffix);
238 out.push_str(line_prefix);
239 }
240
241 v_htmlescape::b_escape(line.as_bytes(), out);
242 }
243
244 if content.ends_with('\n') {
245 out.push_str(line_suffix);
246 tag_open = false;
247 }
248 }
249 HighlightEvent::HighlightStart(highlight) => {
250 write!(
251 out,
252 r#"<span class="highlight {}">"#,
253 HIGHLIGHT_CLASSES[highlight.0]
254 )?;
255 }
256 HighlightEvent::HighlightEnd => {
257 out.push_str("</span>");
258 }
259 }
260 }
261
262 if tag_open {
263 out.push_str(line_suffix);
264 }
265
266 Ok::<_, anyhow::Error>(())
267 })?;
268
269 Ok(())
270}