1use std::collections::BTreeSet;
2use std::sync::Arc;
3
4use crate::completion::{CommandLineParser, CompletionNode, CompletionTree, TokenSpan};
5use nu_ansi_term::Color;
6use reedline::{Highlighter, StyledText};
7use serde::Serialize;
8
9use crate::repl::LineProjection;
10
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub(crate) enum HighlightTokenKind {
18 Plain,
19 CommandValid,
20 ColorLiteral(Color),
21}
22
23#[derive(Debug, Clone, PartialEq, Eq)]
24pub(crate) struct HighlightedSpan {
25 pub start: usize,
26 pub end: usize,
27 pub kind: HighlightTokenKind,
28}
29
30#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
32pub struct HighlightDebugSpan {
33 pub start: usize,
35 pub end: usize,
37 pub text: String,
39 pub kind: String,
41 pub rgb: Option<[u8; 3]>,
43}
44
45pub(crate) type LineProjector = Arc<dyn Fn(&str) -> LineProjection + Send + Sync>;
46
47pub(crate) struct ReplHighlighter {
48 tree: CompletionTree,
49 parser: CommandLineParser,
50 command_color: Color,
51 line_projector: Option<LineProjector>,
52}
53
54impl ReplHighlighter {
55 pub(crate) fn new(
56 tree: CompletionTree,
57 command_color: Color,
58 line_projector: Option<LineProjector>,
59 ) -> Self {
60 Self {
61 tree,
62 parser: CommandLineParser,
63 command_color,
64 line_projector,
65 }
66 }
67
68 pub(crate) fn classify(&self, line: &str) -> Vec<HighlightedSpan> {
69 if line.is_empty() {
70 return Vec::new();
71 }
72
73 let projected = self
74 .line_projector
75 .as_ref()
76 .map(|project| project(line))
77 .unwrap_or_else(|| LineProjection::passthrough(line));
78 let raw_spans = self.parser.tokenize_with_spans(line);
79 if raw_spans.is_empty() {
80 return Vec::new();
81 }
82
83 let mut command_ranges =
84 command_token_ranges(&self.tree.root, &self.parser, &projected.line);
85 if let Some(range) = blanked_help_keyword_range(&raw_spans, &projected.line) {
86 command_ranges.insert(range);
87 }
88
89 raw_spans
90 .into_iter()
91 .map(|span| HighlightedSpan {
92 start: span.start,
93 end: span.end,
94 kind: if command_ranges.contains(&(span.start, span.end)) {
95 HighlightTokenKind::CommandValid
96 } else if let Some(color) = parse_hex_color_token(&span.value) {
97 HighlightTokenKind::ColorLiteral(color)
98 } else {
99 HighlightTokenKind::Plain
100 },
101 })
102 .collect()
103 }
104
105 fn classify_debug(&self, line: &str) -> Vec<HighlightDebugSpan> {
106 self.classify(line)
107 .into_iter()
108 .map(|span| HighlightDebugSpan {
109 start: span.start,
110 end: span.end,
111 text: line[span.start..span.end].to_string(),
112 kind: debug_kind_name(span.kind).to_string(),
113 rgb: debug_kind_rgb(span.kind),
114 })
115 .collect()
116 }
117}
118
119impl Highlighter for ReplHighlighter {
120 fn highlight(&self, line: &str, _cursor: usize) -> StyledText {
121 let mut styled = StyledText::new();
122 if line.is_empty() {
123 return styled;
124 }
125
126 let spans = self.classify(line);
127 if spans.is_empty() {
128 styled.push((nu_ansi_term::Style::new(), line.to_string()));
129 return styled;
130 }
131
132 let mut pos = 0usize;
133 for span in spans {
134 if span.start > pos {
135 styled.push((
136 nu_ansi_term::Style::new(),
137 line[pos..span.start].to_string(),
138 ));
139 }
140
141 let style = match span.kind {
142 HighlightTokenKind::Plain => nu_ansi_term::Style::new(),
143 HighlightTokenKind::CommandValid => {
144 nu_ansi_term::Style::new().fg(self.command_color)
145 }
146 HighlightTokenKind::ColorLiteral(color) => nu_ansi_term::Style::new().fg(color),
147 };
148 styled.push((style, line[span.start..span.end].to_string()));
149 pos = span.end;
150 }
151
152 if pos < line.len() {
153 styled.push((nu_ansi_term::Style::new(), line[pos..].to_string()));
154 }
155
156 styled
157 }
158}
159
160pub fn debug_highlight(
162 tree: &CompletionTree,
163 line: &str,
164 command_color: Color,
165 line_projector: Option<LineProjector>,
166) -> Vec<HighlightDebugSpan> {
167 ReplHighlighter::new(tree.clone(), command_color, line_projector).classify_debug(line)
168}
169
170fn command_token_ranges(
171 root: &CompletionNode,
172 parser: &CommandLineParser,
173 projected_line: &str,
174) -> BTreeSet<(usize, usize)> {
175 let mut ranges = BTreeSet::new();
176 let spans = parser.tokenize_with_spans(projected_line);
177 if spans.is_empty() {
178 return ranges;
179 }
180
181 let mut node = root;
182 for span in spans {
183 let token = span.value.as_str();
184 if token.is_empty() || token == "|" || token.starts_with('-') {
185 break;
186 }
187
188 let Some(child) = node.children.get(token) else {
189 break;
190 };
191
192 ranges.insert((span.start, span.end));
193 node = child;
194 }
195
196 ranges
197}
198
199fn blanked_help_keyword_range(
202 raw_spans: &[TokenSpan],
203 projected_line: &str,
204) -> Option<(usize, usize)> {
205 raw_spans
206 .iter()
207 .find(|span| {
208 span.value == "help"
209 && projected_line
210 .get(span.start..span.end)
211 .is_some_and(|segment| segment.trim().is_empty())
212 })
213 .map(|span| (span.start, span.end))
214}
215
216fn parse_hex_color_token(token: &str) -> Option<Color> {
217 let normalized = token.trim();
218 let hex = normalized.strip_prefix('#')?;
219 if hex.len() == 6 {
220 let r = u8::from_str_radix(&hex[0..2], 16).ok()?;
221 let g = u8::from_str_radix(&hex[2..4], 16).ok()?;
222 let b = u8::from_str_radix(&hex[4..6], 16).ok()?;
223 return Some(Color::Rgb(r, g, b));
224 }
225 if hex.len() == 3 {
226 let r = u8::from_str_radix(&hex[0..1], 16).ok()?;
227 let g = u8::from_str_radix(&hex[1..2], 16).ok()?;
228 let b = u8::from_str_radix(&hex[2..3], 16).ok()?;
229 return Some(Color::Rgb(
230 r.saturating_mul(17),
231 g.saturating_mul(17),
232 b.saturating_mul(17),
233 ));
234 }
235 None
236}
237
238fn debug_kind_name(kind: HighlightTokenKind) -> &'static str {
239 match kind {
240 HighlightTokenKind::Plain => "plain",
241 HighlightTokenKind::CommandValid => "command_valid",
242 HighlightTokenKind::ColorLiteral(_) => "color_literal",
243 }
244}
245
246fn debug_kind_rgb(kind: HighlightTokenKind) -> Option<[u8; 3]> {
247 let color = match kind {
248 HighlightTokenKind::ColorLiteral(color) => color,
249 _ => return None,
250 };
251
252 let rgb = match color {
253 Color::Black => [0, 0, 0],
254 Color::DarkGray => [128, 128, 128],
255 Color::Red => [128, 0, 0],
256 Color::Green => [0, 128, 0],
257 Color::Yellow => [128, 128, 0],
258 Color::Blue => [0, 0, 128],
259 Color::Purple => [128, 0, 128],
260 Color::Magenta => [128, 0, 128],
261 Color::Cyan => [0, 128, 128],
262 Color::White => [192, 192, 192],
263 Color::Fixed(_) => return None,
264 Color::LightRed => [255, 0, 0],
265 Color::LightGreen => [0, 255, 0],
266 Color::LightYellow => [255, 255, 0],
267 Color::LightBlue => [0, 0, 255],
268 Color::LightPurple => [255, 0, 255],
269 Color::LightMagenta => [255, 0, 255],
270 Color::LightCyan => [0, 255, 255],
271 Color::LightGray => [255, 255, 255],
272 Color::Rgb(r, g, b) => [r, g, b],
273 Color::Default => return None,
274 };
275 Some(rgb)
276}
277
278#[cfg(test)]
279mod tests {
280 use super::{ReplHighlighter, debug_highlight};
281 use crate::completion::{CompletionNode, CompletionTree};
282 use crate::repl::LineProjection;
283 use nu_ansi_term::Color;
284 use reedline::Highlighter;
285 use std::sync::Arc;
286
287 fn token_styles(styled: &StyledText) -> Vec<(String, Option<Color>)> {
288 styled
289 .buffer
290 .iter()
291 .filter_map(|(style, text)| {
292 if text.chars().all(|ch| ch.is_whitespace()) {
293 None
294 } else {
295 Some((text.clone(), style.foreground))
296 }
297 })
298 .collect()
299 }
300
301 use reedline::StyledText;
302
303 fn completion_tree_with_config_show() -> CompletionTree {
304 let mut config = CompletionNode::default();
305 config
306 .children
307 .insert("show".to_string(), CompletionNode::default());
308 CompletionTree {
309 root: CompletionNode::default().with_child("config", config),
310 ..CompletionTree::default()
311 }
312 }
313
314 #[test]
315 fn colors_full_command_chain_only_unit() {
316 let tree = completion_tree_with_config_show();
317 let highlighter = ReplHighlighter::new(tree, Color::Green, None);
318
319 let tokens = token_styles(&highlighter.highlight("config show", 0));
320 assert_eq!(
321 tokens,
322 vec![
323 ("config".to_string(), Some(Color::Green)),
324 ("show".to_string(), Some(Color::Green)),
325 ]
326 );
327 }
328
329 #[test]
330 fn skips_partial_subcommand_and_flags_unit() {
331 let tree = completion_tree_with_config_show();
332 let highlighter = ReplHighlighter::new(tree, Color::Green, None);
333
334 let tokens = token_styles(&highlighter.highlight("config sho", 0));
335 assert_eq!(
336 tokens,
337 vec![
338 ("config".to_string(), Some(Color::Green)),
339 ("sho".to_string(), None),
340 ]
341 );
342
343 let tokens = token_styles(&highlighter.highlight("config --flag", 0));
344 assert_eq!(
345 tokens,
346 vec![
347 ("config".to_string(), Some(Color::Green)),
348 ("--flag".to_string(), None),
349 ]
350 );
351 }
352
353 #[test]
354 fn colors_help_alias_keyword_and_target_unit() {
355 let tree = CompletionTree {
356 root: CompletionNode::default().with_child("history", CompletionNode::default()),
357 ..CompletionTree::default()
358 };
359 let projector =
360 Arc::new(|line: &str| LineProjection::passthrough(line.replacen("help", " ", 1)));
361 let highlighter = ReplHighlighter::new(tree, Color::Green, Some(projector));
362
363 let tokens = token_styles(&highlighter.highlight("help history", 0));
364 assert_eq!(
365 tokens,
366 vec![
367 ("help".to_string(), Some(Color::Green)),
368 ("history".to_string(), Some(Color::Green)),
369 ]
370 );
371
372 let tokens = token_styles(&highlighter.highlight("help his", 0));
373 assert_eq!(
374 tokens,
375 vec![
376 ("help".to_string(), Some(Color::Green)),
377 ("his".to_string(), None),
378 ]
379 );
380 }
381
382 #[test]
383 fn highlights_hex_color_literals_unit() {
384 let highlighter = ReplHighlighter::new(CompletionTree::default(), Color::Green, None);
385 let spans = debug_highlight(&CompletionTree::default(), "#ff00cc", Color::Green, None);
386 assert_eq!(spans.len(), 1);
387 assert_eq!(spans[0].kind, "color_literal");
388 assert_eq!(spans[0].rgb, Some([255, 0, 204]));
389 let tokens = token_styles(&highlighter.highlight("#ff00cc", 0));
390 assert_eq!(
391 tokens,
392 vec![("#ff00cc".to_string(), Some(Color::Rgb(255, 0, 204)))]
393 );
394 }
395
396 #[test]
397 fn debug_spans_preserve_help_alias_ranges_unit() {
398 let tree = CompletionTree {
399 root: CompletionNode::default().with_child("history", CompletionNode::default()),
400 ..CompletionTree::default()
401 };
402 let projector =
403 Arc::new(|line: &str| LineProjection::passthrough(line.replacen("help", " ", 1)));
404 let spans = debug_highlight(&tree, "help history -", Color::Green, Some(projector));
405
406 assert_eq!(
407 spans
408 .into_iter()
409 .filter(|span| span.kind == "command_valid")
410 .map(|span| (span.start, span.end, span.text))
411 .collect::<Vec<_>>(),
412 vec![(0, 4, "help".to_string()), (5, 12, "history".to_string())]
413 );
414 }
415
416 #[test]
417 fn three_digit_hex_and_invalid_tokens_cover_debug_paths_unit() {
418 let spans = debug_highlight(&CompletionTree::default(), "#0af", Color::Green, None);
419 assert_eq!(spans[0].rgb, Some([0, 170, 255]));
420
421 let highlighter = ReplHighlighter::new(CompletionTree::default(), Color::Green, None);
422 let tokens = token_styles(&highlighter.highlight("unknown #nope", 0));
423 assert_eq!(
424 tokens,
425 vec![("unknown".to_string(), None), ("#nope".to_string(), None),]
426 );
427 }
428}