Skip to main content

vtcode_tui/ui/markdown/
code_blocks.rs

1use super::parsing::{flush_current_line, push_blank_line};
2use super::{CODE_LINE_NUMBER_MIN_WIDTH, MarkdownLine, MarkdownSegment, RenderMarkdownOptions};
3use crate::config::loader::SyntaxHighlightingConfig;
4use crate::ui::syntax_highlight;
5use crate::ui::theme::ThemeStyles;
6use crate::utils::diff_styles::DiffColorPalette;
7use anstyle::{Effects, Style};
8use pulldown_cmark::{Event, Options, Parser, Tag, TagEnd};
9use syntect::util::LinesWithEndings;
10use vtcode_commons::diff_paths::{
11    format_start_only_hunk_header, is_diff_addition_line, is_diff_deletion_line,
12    is_diff_header_line, is_diff_new_file_marker_line, looks_like_diff_content,
13    parse_diff_git_path, parse_diff_marker_path,
14};
15
16const DIFF_SUMMARY_PREFIX: &str = "• Diff ";
17
18#[derive(Clone, Debug)]
19pub(crate) struct CodeBlockState {
20    pub(crate) language: Option<String>,
21    pub(crate) buffer: String,
22}
23
24pub(crate) struct CodeBlockRenderEnv<'a> {
25    pub(crate) lines: &'a mut Vec<MarkdownLine>,
26    pub(crate) current_line: &'a mut MarkdownLine,
27    pub(crate) blockquote_depth: usize,
28    pub(crate) list_continuation_prefix: &'a str,
29    pub(crate) pending_list_prefix: &'a mut Option<String>,
30    pub(crate) base_style: Style,
31    pub(crate) theme_styles: &'a ThemeStyles,
32    pub(crate) highlight_config: Option<&'a SyntaxHighlightingConfig>,
33    pub(crate) render_options: RenderMarkdownOptions,
34}
35
36pub(crate) fn handle_code_block_event(
37    event: &Event<'_>,
38    code_block: &mut Option<CodeBlockState>,
39    env: &mut CodeBlockRenderEnv<'_>,
40) -> bool {
41    if code_block.is_none() {
42        return false;
43    }
44
45    match event {
46        Event::Text(text) => {
47            if let Some(state) = code_block.as_mut() {
48                state.buffer.push_str(text);
49            }
50            true
51        }
52        Event::End(TagEnd::CodeBlock) => {
53            finalize_code_block(env, code_block, true, true);
54            true
55        }
56        _ => false,
57    }
58}
59
60pub(crate) fn finalize_unclosed_code_block(
61    code_block: &mut Option<CodeBlockState>,
62    env: &mut CodeBlockRenderEnv<'_>,
63) {
64    finalize_code_block(env, code_block, false, false);
65}
66
67fn finalize_code_block(
68    env: &mut CodeBlockRenderEnv<'_>,
69    code_block: &mut Option<CodeBlockState>,
70    allow_table_reparse: bool,
71    append_trailing_blank_line: bool,
72) {
73    flush_current_line(
74        env.lines,
75        env.current_line,
76        env.blockquote_depth,
77        env.list_continuation_prefix,
78        env.pending_list_prefix,
79        env.base_style,
80    );
81    if let Some(state) = code_block.take() {
82        let rendered = render_code_block_state(&state, env, allow_table_reparse);
83        env.lines.extend(rendered);
84        if append_trailing_blank_line {
85            push_blank_line(env.lines);
86        }
87    }
88}
89
90fn render_code_block_state(
91    state: &CodeBlockState,
92    env: &CodeBlockRenderEnv<'_>,
93    allow_table_reparse: bool,
94) -> Vec<MarkdownLine> {
95    if allow_table_reparse
96        && !env.render_options.disable_code_block_table_reparse
97        && code_block_contains_table(&state.buffer, state.language.as_deref())
98    {
99        return render_markdown_code_block_table(
100            &state.buffer,
101            env.base_style,
102            env.theme_styles,
103            env.highlight_config,
104            env.render_options,
105        );
106    }
107
108    let prefix = build_prefix_segments(
109        env.blockquote_depth,
110        env.list_continuation_prefix,
111        env.base_style,
112    );
113    highlight_code_block(
114        &state.buffer,
115        state.language.as_deref(),
116        env.highlight_config,
117        env.theme_styles,
118        env.base_style,
119        &prefix,
120        env.render_options.preserve_code_indentation,
121    )
122}
123
124fn render_markdown_code_block_table(
125    source: &str,
126    base_style: Style,
127    theme_styles: &ThemeStyles,
128    highlight_config: Option<&SyntaxHighlightingConfig>,
129    render_options: RenderMarkdownOptions,
130) -> Vec<MarkdownLine> {
131    let mut nested_options = render_options;
132    nested_options.disable_code_block_table_reparse = true;
133    super::render_markdown_to_lines_with_options(
134        source,
135        base_style,
136        theme_styles,
137        highlight_config,
138        nested_options,
139    )
140}
141
142fn build_prefix_segments(
143    blockquote_depth: usize,
144    list_continuation_prefix: &str,
145    base_style: Style,
146) -> Vec<MarkdownSegment> {
147    let mut segments =
148        Vec::with_capacity(blockquote_depth + usize::from(!list_continuation_prefix.is_empty()));
149    for _ in 0..blockquote_depth {
150        segments.push(MarkdownSegment::new(base_style.dimmed().italic(), "│ "));
151    }
152    if !list_continuation_prefix.is_empty() {
153        segments.push(MarkdownSegment::new(base_style, list_continuation_prefix));
154    }
155    segments
156}
157
158fn highlight_code_block(
159    code: &str,
160    language: Option<&str>,
161    highlight_config: Option<&SyntaxHighlightingConfig>,
162    theme_styles: &ThemeStyles,
163    base_style: Style,
164    prefix_segments: &[MarkdownSegment],
165    preserve_code_indentation: bool,
166) -> Vec<MarkdownLine> {
167    let mut lines = Vec::new();
168
169    let normalized_code = normalize_code_indentation(code, language, preserve_code_indentation);
170    let code_to_display = &normalized_code;
171    if is_diff_language(language)
172        || (language.is_none() && looks_like_diff_content(code_to_display))
173    {
174        return render_diff_code_block(code_to_display, theme_styles, base_style, prefix_segments);
175    }
176    let use_line_numbers =
177        language.is_some_and(|lang| !lang.trim().is_empty()) && !is_diff_language(language);
178
179    if let Some(config) = highlight_config.filter(|cfg| cfg.enabled)
180        && let Some(highlighted) = try_highlight(code_to_display, language, config)
181    {
182        let line_count = highlighted.len();
183        let number_width = line_number_width(line_count);
184        for (index, segments) in highlighted.into_iter().enumerate() {
185            let mut line = code_line_with_prefix(
186                prefix_segments,
187                theme_styles,
188                base_style,
189                use_line_numbers.then_some((index + 1, number_width)),
190            );
191            for (style, text) in segments {
192                line.push_segment(style, &text);
193            }
194            lines.push(line);
195        }
196        return lines;
197    }
198
199    let mut line_number = 1usize;
200    let mut line_count = LinesWithEndings::from(code_to_display).count();
201    if code_to_display.ends_with('\n') {
202        line_count = line_count.saturating_add(1);
203    }
204    let number_width = line_number_width(line_count);
205
206    for raw_line in LinesWithEndings::from(code_to_display) {
207        let trimmed = raw_line.trim_end_matches('\n');
208        let mut line = code_line_with_prefix(
209            prefix_segments,
210            theme_styles,
211            base_style,
212            use_line_numbers.then_some((line_number, number_width)),
213        );
214        if !trimmed.is_empty() {
215            line.push_segment(code_block_style(theme_styles, base_style), trimmed);
216        }
217        lines.push(line);
218        line_number = line_number.saturating_add(1);
219    }
220
221    if code_to_display.ends_with('\n') {
222        let line = code_line_with_prefix(
223            prefix_segments,
224            theme_styles,
225            base_style,
226            use_line_numbers.then_some((line_number, number_width)),
227        );
228        lines.push(line);
229    }
230
231    lines
232}
233
234pub(crate) fn normalize_diff_lines(code: &str) -> Vec<String> {
235    #[derive(Default)]
236    struct DiffBlock {
237        header: String,
238        path: String,
239        lines: Vec<String>,
240        additions: usize,
241        deletions: usize,
242    }
243
244    let mut preface = Vec::new();
245    let mut blocks = Vec::new();
246    let mut current: Option<DiffBlock> = None;
247    let mut fallback_additions = 0usize;
248    let mut fallback_deletions = 0usize;
249    let mut fallback_path: Option<String> = None;
250    let mut summary_insert_index: Option<usize> = None;
251
252    for line in code.lines() {
253        if fallback_path.is_none() {
254            fallback_path = parse_diff_marker_path(line);
255        }
256        if summary_insert_index.is_none() && is_diff_new_file_marker_line(line.trim_start()) {
257            summary_insert_index = Some(preface.len());
258        }
259        bump_diff_counters(line, &mut fallback_additions, &mut fallback_deletions);
260
261        if let Some(path) = parse_diff_git_path(line) {
262            if let Some(block) = current.take() {
263                blocks.push(block);
264            }
265            current = Some(DiffBlock {
266                header: line.to_string(),
267                path,
268                lines: Vec::new(),
269                additions: 0,
270                deletions: 0,
271            });
272            continue;
273        }
274
275        let rewritten = rewrite_diff_line(line);
276        if let Some(block) = current.as_mut() {
277            bump_diff_counters(line, &mut block.additions, &mut block.deletions);
278            block.lines.push(rewritten);
279        } else {
280            preface.push(rewritten);
281        }
282    }
283
284    if let Some(block) = current {
285        blocks.push(block);
286    }
287
288    if blocks.is_empty() {
289        let path = fallback_path.unwrap_or_else(|| "file".to_string());
290        let summary = format_diff_summary(path.as_str(), fallback_additions, fallback_deletions);
291
292        let mut output = Vec::with_capacity(preface.len() + 1);
293        if let Some(idx) = summary_insert_index {
294            output.extend(preface[..=idx].iter().cloned());
295            output.push(summary);
296            output.extend(preface[idx + 1..].iter().cloned());
297        } else {
298            output.push(summary);
299            output.extend(preface);
300        }
301        return output;
302    }
303
304    let mut output = Vec::new();
305    output.extend(preface);
306    for block in blocks {
307        output.push(block.header);
308        output.push(format_diff_summary(
309            block.path.as_str(),
310            block.additions,
311            block.deletions,
312        ));
313        output.extend(block.lines);
314    }
315    output
316}
317
318fn render_diff_code_block(
319    code: &str,
320    theme_styles: &ThemeStyles,
321    base_style: Style,
322    prefix_segments: &[MarkdownSegment],
323) -> Vec<MarkdownLine> {
324    let mut lines = Vec::new();
325    let palette = DiffColorPalette::default();
326    let context_style = code_block_style(theme_styles, base_style);
327    let header_style = palette.header_style();
328    let added_style = palette.added_style();
329    let removed_style = palette.removed_style();
330
331    for line in normalize_diff_lines(code) {
332        let trimmed = line.trim_end_matches('\n');
333        let trimmed_start = trimmed.trim_start();
334        if let Some((path, additions, deletions)) = parse_diff_summary_line(trimmed_start) {
335            let leading_len = trimmed.len().saturating_sub(trimmed_start.len());
336            let leading = &trimmed[..leading_len];
337            let mut line = prefixed_line(prefix_segments);
338            if !leading.is_empty() {
339                line.push_segment(context_style, leading);
340            }
341            line.push_segment(context_style, &format!("{DIFF_SUMMARY_PREFIX}{path} ("));
342            line.push_segment(added_style, &format!("+{additions}"));
343            line.push_segment(context_style, " ");
344            line.push_segment(removed_style, &format!("-{deletions}"));
345            line.push_segment(context_style, ")");
346            lines.push(line);
347            continue;
348        }
349        let style = if trimmed.is_empty() {
350            context_style
351        } else if is_diff_header_line(trimmed_start) {
352            header_style
353        } else if is_diff_addition_line(trimmed_start) {
354            added_style
355        } else if is_diff_deletion_line(trimmed_start) {
356            removed_style
357        } else {
358            context_style
359        };
360
361        let mut line = prefixed_line(prefix_segments);
362        if !trimmed.is_empty() {
363            line.push_segment(style, trimmed);
364        }
365        lines.push(line);
366    }
367
368    if code.ends_with('\n') {
369        lines.push(prefixed_line(prefix_segments));
370    }
371
372    lines
373}
374
375fn parse_diff_summary_line(line: &str) -> Option<(&str, usize, usize)> {
376    let summary = line.strip_prefix(DIFF_SUMMARY_PREFIX)?;
377    let (path, counts) = summary.rsplit_once(" (")?;
378    let counts = counts.strip_suffix(')')?;
379    let mut parts = counts.split_whitespace();
380    let additions = parts.next()?.strip_prefix('+')?.parse().ok()?;
381    let deletions = parts.next()?.strip_prefix('-')?.parse().ok()?;
382    Some((path, additions, deletions))
383}
384
385fn format_diff_summary(path: &str, additions: usize, deletions: usize) -> String {
386    format!("{DIFF_SUMMARY_PREFIX}{path} (+{additions} -{deletions})")
387}
388
389fn append_prefix_segments(line: &mut MarkdownLine, prefix_segments: &[MarkdownSegment]) {
390    for segment in prefix_segments {
391        line.push_segment(segment.style, &segment.text);
392    }
393}
394
395fn prefixed_line(prefix_segments: &[MarkdownSegment]) -> MarkdownLine {
396    let mut line = MarkdownLine::default();
397    append_prefix_segments(&mut line, prefix_segments);
398    line
399}
400
401fn append_code_line_prefix(
402    line: &mut MarkdownLine,
403    prefix_segments: &[MarkdownSegment],
404    theme_styles: &ThemeStyles,
405    base_style: Style,
406    line_number: Option<(usize, usize)>,
407) {
408    append_prefix_segments(line, prefix_segments);
409    let Some((line_number, width)) = line_number else {
410        return;
411    };
412
413    let number_text = format!("{line_number:>width$}  ");
414    let number_style = if base_style == theme_styles.tool_output {
415        theme_styles.tool_detail.dimmed()
416    } else {
417        base_style.dimmed()
418    };
419    line.push_segment(number_style, &number_text);
420}
421
422fn code_line_with_prefix(
423    prefix_segments: &[MarkdownSegment],
424    theme_styles: &ThemeStyles,
425    base_style: Style,
426    line_number: Option<(usize, usize)>,
427) -> MarkdownLine {
428    let mut line = MarkdownLine::default();
429    append_code_line_prefix(
430        &mut line,
431        prefix_segments,
432        theme_styles,
433        base_style,
434        line_number,
435    );
436    line
437}
438
439fn line_number_width(line_count: usize) -> usize {
440    let digits = line_count.max(1).to_string().len();
441    digits.max(CODE_LINE_NUMBER_MIN_WIDTH)
442}
443
444fn code_block_contains_table(content: &str, language: Option<&str>) -> bool {
445    if let Some(lang) = language {
446        let lang_lower = lang.to_ascii_lowercase();
447        if !matches!(
448            lang_lower.as_str(),
449            "markdown" | "md" | "text" | "txt" | "plaintext" | "plain"
450        ) {
451            return false;
452        }
453    }
454
455    let trimmed = content.trim();
456    if trimmed.is_empty() {
457        return false;
458    }
459
460    let mut has_pipe_line = false;
461    let mut has_separator = false;
462    for line in trimmed.lines().take(4) {
463        let line = line.trim();
464        if line.contains('|') {
465            has_pipe_line = true;
466        }
467        if line.starts_with('|') && line.chars().all(|c| matches!(c, '|' | '-' | ':' | ' ')) {
468            has_separator = true;
469        }
470    }
471    if !has_pipe_line || !has_separator {
472        return false;
473    }
474
475    let options = Options::ENABLE_TABLES;
476    let parser = Parser::new_ext(trimmed, options);
477    for event in parser {
478        match event {
479            Event::Start(Tag::Table(_)) => return true,
480            Event::Start(Tag::Paragraph) | Event::Text(_) | Event::SoftBreak => continue,
481            _ => return false,
482        }
483    }
484    false
485}
486
487fn rewrite_diff_line(line: &str) -> String {
488    format_start_only_hunk_header(line).unwrap_or_else(|| line.to_string())
489}
490
491fn bump_diff_counters(line: &str, additions: &mut usize, deletions: &mut usize) {
492    let trimmed = line.trim_start();
493    if is_diff_addition_line(trimmed) {
494        *additions += 1;
495    } else if is_diff_deletion_line(trimmed) {
496        *deletions += 1;
497    }
498}
499
500fn is_diff_language(language: Option<&str>) -> bool {
501    language.is_some_and(|lang| {
502        matches!(
503            lang.to_ascii_lowercase().as_str(),
504            "diff" | "patch" | "udiff" | "git"
505        )
506    })
507}
508
509fn code_block_style(theme_styles: &ThemeStyles, base_style: Style) -> Style {
510    let base_fg = base_style.get_fg_color();
511    let theme_fg = theme_styles.output.get_fg_color();
512    let fg = if base_style.get_effects().contains(Effects::DIMMED) {
513        base_fg.or(theme_fg)
514    } else {
515        theme_fg.or(base_fg)
516    };
517    let mut style = base_style;
518    if let Some(color) = fg {
519        style = style.fg_color(Some(color));
520    }
521    style
522}
523
524pub(crate) fn normalize_code_indentation(
525    code: &str,
526    language: Option<&str>,
527    preserve_indentation: bool,
528) -> String {
529    if preserve_indentation {
530        return code.to_string();
531    }
532    let has_language_hint = language.is_some_and(|hint| {
533        matches!(
534            hint.to_lowercase().as_str(),
535            "rust"
536                | "rs"
537                | "python"
538                | "py"
539                | "javascript"
540                | "js"
541                | "jsx"
542                | "typescript"
543                | "ts"
544                | "tsx"
545                | "go"
546                | "golang"
547                | "java"
548                | "cpp"
549                | "c"
550                | "php"
551                | "html"
552                | "css"
553                | "sql"
554                | "csharp"
555                | "bash"
556                | "sh"
557                | "swift"
558        )
559    });
560
561    if !has_language_hint && language.is_some() {
562        return code.to_string();
563    }
564
565    let lines: Vec<&str> = code.lines().collect();
566    let min_indent = lines
567        .iter()
568        .filter(|line| !line.trim().is_empty())
569        .map(|line| &line[..line.len() - line.trim_start().len()])
570        .reduce(|acc, p| {
571            let mut len = 0;
572            for (c1, c2) in acc.chars().zip(p.chars()) {
573                if c1 != c2 {
574                    break;
575                }
576                len += c1.len_utf8();
577            }
578            &acc[..len]
579        })
580        .map(|s| s.len())
581        .unwrap_or(0);
582
583    let normalized = lines
584        .iter()
585        .map(|line| {
586            if line.trim().is_empty() {
587                line
588            } else if line.len() >= min_indent {
589                &line[min_indent..]
590            } else {
591                line
592            }
593        })
594        .collect::<Vec<_>>()
595        .join("\n");
596
597    if code.ends_with('\n') {
598        format!("{normalized}\n")
599    } else {
600        normalized
601    }
602}
603
604pub fn highlight_line_for_diff(line: &str, language: Option<&str>) -> Option<Vec<(Style, String)>> {
605    syntax_highlight::highlight_line_to_anstyle_segments(
606        line,
607        language,
608        syntax_highlight::get_active_syntax_theme(),
609        true,
610    )
611    .map(|segments| {
612        segments
613            .into_iter()
614            .map(|(style, text)| {
615                let fg = style.get_fg_color().map(|c| match c {
616                    anstyle::Color::Rgb(rgb) => {
617                        let brighten = |v: u8| (v as u16 * 120 / 100).min(255) as u8;
618                        anstyle::Color::Rgb(anstyle::RgbColor(
619                            brighten(rgb.0),
620                            brighten(rgb.1),
621                            brighten(rgb.2),
622                        ))
623                    }
624                    anstyle::Color::Ansi(ansi) => match ansi {
625                        anstyle::AnsiColor::Black => {
626                            anstyle::Color::Ansi(anstyle::AnsiColor::BrightWhite)
627                        }
628                        anstyle::AnsiColor::Red => {
629                            anstyle::Color::Ansi(anstyle::AnsiColor::BrightRed)
630                        }
631                        anstyle::AnsiColor::Green => {
632                            anstyle::Color::Ansi(anstyle::AnsiColor::BrightGreen)
633                        }
634                        anstyle::AnsiColor::Yellow => {
635                            anstyle::Color::Ansi(anstyle::AnsiColor::BrightYellow)
636                        }
637                        anstyle::AnsiColor::Blue => {
638                            anstyle::Color::Ansi(anstyle::AnsiColor::BrightBlue)
639                        }
640                        anstyle::AnsiColor::Magenta => {
641                            anstyle::Color::Ansi(anstyle::AnsiColor::BrightMagenta)
642                        }
643                        anstyle::AnsiColor::Cyan => {
644                            anstyle::Color::Ansi(anstyle::AnsiColor::BrightCyan)
645                        }
646                        anstyle::AnsiColor::White => {
647                            anstyle::Color::Ansi(anstyle::AnsiColor::BrightWhite)
648                        }
649                        other => anstyle::Color::Ansi(other),
650                    },
651                    other => other,
652                });
653                let bg = style.get_bg_color();
654                let new_style = style.fg_color(fg).bg_color(bg);
655                (new_style, text)
656            })
657            .collect()
658    })
659}
660
661fn try_highlight(
662    code: &str,
663    language: Option<&str>,
664    config: &SyntaxHighlightingConfig,
665) -> Option<Vec<Vec<(Style, String)>>> {
666    let max_bytes = config.max_file_size_mb.saturating_mul(1024 * 1024);
667    if max_bytes > 0 && code.len() > max_bytes {
668        return None;
669    }
670
671    if let Some(lang) = language
672        && !config.enabled_languages.is_empty()
673    {
674        let direct_match = config
675            .enabled_languages
676            .iter()
677            .any(|entry| entry.eq_ignore_ascii_case(lang));
678        if !direct_match {
679            let syntax_ref = syntax_highlight::find_syntax_by_token(lang);
680            let resolved_match = config
681                .enabled_languages
682                .iter()
683                .any(|entry| entry.eq_ignore_ascii_case(&syntax_ref.name));
684            if !resolved_match {
685                return None;
686            }
687        }
688    }
689
690    let rendered = syntax_highlight::highlight_code_to_anstyle_line_segments(
691        code,
692        language,
693        &config.theme,
694        true,
695    );
696
697    Some(rendered)
698}
699
700#[derive(Clone, Debug)]
701pub struct HighlightedSegment {
702    pub style: Style,
703    pub text: String,
704}
705
706pub fn highlight_code_to_segments(
707    code: &str,
708    language: Option<&str>,
709    theme_name: &str,
710) -> Vec<Vec<HighlightedSegment>> {
711    syntax_highlight::highlight_code_to_anstyle_line_segments(code, language, theme_name, true)
712        .into_iter()
713        .map(|segments| {
714            segments
715                .into_iter()
716                .map(|(style, text)| HighlightedSegment { style, text })
717                .collect()
718        })
719        .collect()
720}
721
722pub fn highlight_code_to_ansi(code: &str, language: Option<&str>, theme_name: &str) -> Vec<String> {
723    let segments = highlight_code_to_segments(code, language, theme_name);
724    segments
725        .into_iter()
726        .map(|line_segments| {
727            let mut ansi_line = String::new();
728            for seg in line_segments {
729                let rendered = seg.style.render();
730                ansi_line.push_str(&format!(
731                    "{rendered}{text}{reset}",
732                    text = seg.text,
733                    reset = anstyle::Reset
734                ));
735            }
736            ansi_line
737        })
738        .collect()
739}