rolldown_error 1.3.2

use std::{fmt::Display, fmt::Write as _, ops::Range};

use arcstr::ArcStr;
use ariadne::{Config, Label, Report, ReportBuilder, ReportKind, Span, sources};
use rustc_hash::FxHashMap;

use crate::utils::{ByteLocator, is_context_too_long};

use super::Severity;

#[derive(Debug, Clone, PartialEq, PartialOrd, Hash, Eq)]
pub struct DiagnosticFileId(ArcStr);

impl Display for DiagnosticFileId {
  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
    self.0.fmt(f)
  }
}

impl From<ArcStr> for DiagnosticFileId {
  fn from(value: ArcStr) -> Self {
    Self(value)
  }
}

#[derive(Debug, Clone)]
pub struct RolldownLabelSpan(DiagnosticFileId, Range<usize>);

/// Diagnostics reach users through stderr (rolldown CLI and Vite print errors and warnings there).
pub fn stderr_supports_color() -> bool {
  supports_color::on_cached(supports_color::Stream::Stderr).is_some()
}

impl ariadne::Span for RolldownLabelSpan {
  type SourceId = DiagnosticFileId;

  fn source(&self) -> &Self::SourceId {
    &self.0
  }

  fn start(&self) -> usize {
    self.1.start
  }

  fn end(&self) -> usize {
    self.1.end
  }
}

#[derive(Debug, Clone)]
pub struct Diagnostic {
  pub(crate) kind: String,
  pub(crate) title: String,
  pub(crate) files: FxHashMap</* filename */ DiagnosticFileId, /* file content */ ArcStr>,
  pub(crate) labels: Vec<Label<RolldownLabelSpan>>,
  pub(crate) helps: Vec<String>,
  pub(crate) note: Option<String>,
  pub(crate) severity: Severity,
}

type AriadneReportBuilder = ReportBuilder<'static, RolldownLabelSpan>;
type AriadneReport = Report<'static, RolldownLabelSpan>;

impl Diagnostic {
  pub(crate) fn new(kind: String, summary: String, severity: Severity) -> Self {
    Self {
      kind,
      title: summary,
      files: FxHashMap::default(),
      labels: Vec::default(),
      helps: Vec::default(),
      note: None,
      severity,
    }
  }

  pub(crate) fn add_file(
    &mut self,
    filename: impl Into<ArcStr>,
    content: impl Into<ArcStr>,
  ) -> DiagnosticFileId {
    let filename = DiagnosticFileId::from(filename.into());
    self.files.entry(filename.clone()).or_insert_with(|| content.into());
    filename
  }

  #[inline]
  pub(crate) fn add_help(&mut self, message: String) -> &mut Self {
    self.helps.push(message);
    self
  }

  #[inline]
  pub(crate) fn add_note(&mut self, message: String) -> &mut Self {
    self.note = Some(message);
    self
  }

  pub(crate) fn add_label(
    &mut self,
    file_id: &DiagnosticFileId,
    range: impl Into<Range<u32>>,
    message: String,
  ) -> &mut Self {
    let range = range.into();
    let range = range.start as usize..range.end as usize;
    let label =
      Label::new(RolldownLabelSpan(file_id.0.clone().into(), range)).with_message(message);
    self.labels.push(label);
    self
  }

  fn init_report_builder(&self) -> AriadneReportBuilder {
    let mut message = self.title.clone();
    let mut builder = AriadneReport::build(
      match self.severity {
        Severity::Info => ReportKind::Advice,
        Severity::Error => ReportKind::Error,
        Severity::Warning => ReportKind::Warning,
      },
      RolldownLabelSpan(ArcStr::default().into(), 0..0),
    )
    .with_code(self.kind.clone());

    for label in self.labels.clone() {
      if is_context_too_long(&label, &self.files) {
        let span = label.span();
        let _ = write!(
          message,
          "\n - {} in {} at {:?}",
          label.display_info().msg().unwrap_or_default(),
          span.source(),
          span.start()..span.end()
        );
      } else {
        builder = builder.with_label(label);
      }
    }

    builder.with_helps(self.helps.clone());

    if let Some(note) = &self.note {
      builder = builder.with_note(note);
    }

    builder = builder.with_message(message);

    builder
  }

  pub fn convert_to_string(&self, color: bool) -> String {
    let mut cache = sources(self.files.clone());
    self.convert_to_string_with_cache(color, &mut cache)
  }

  /// Like [`Self::convert_to_string`], but renders against a caller-owned ariadne
  /// source cache so the line-indexed `Source` for each file is built once and
  /// reused across many diagnostics instead of rebuilt per render (see #9748).
  fn convert_to_string_with_cache(
    &self,
    color: bool,
    cache: &mut impl ariadne::Cache<DiagnosticFileId>,
  ) -> String {
    let builder = self.init_report_builder();
    let mut output = Vec::new();
    let result = builder
      .with_config(
        Config::default()
          .with_color(color)
          .with_index_type(ariadne::IndexType::Byte)
          .with_severity_prefix(false),
      )
      .finish()
      .write_for_stdout(&mut *cache, &mut output);
    match result {
      Ok(()) => String::from_utf8_lossy(&output).into_owned(),
      Err(_) => format!("[{}] {}", self.kind, self.title),
    }
  }

  /// Colored only when stderr supports it, so `NO_COLOR`, `FORCE_COLOR` and piped output are respected.
  pub fn to_stderr_string(&self) -> String {
    self.convert_to_string(stderr_supports_color())
  }

  /// Get the primary location information from the first label if available
  /// Returns (file_path, line, column, utf16_position)
  pub fn get_primary_location(&self) -> Option<(String, usize, usize, usize)> {
    let first_label = self.labels.first()?;
    let span = first_label.span();
    let source = self.files.get(span.source())?;

    // A start inside a char counts that whole char, and a start past the end counts the file.
    let prefix = &source[..source.ceil_char_boundary(span.start())];
    let line_start = prefix.rfind('\n').map_or(0, |i| i + 1);
    let line = 1 + memchr::memchr_iter(b'\n', &prefix.as_bytes()[..line_start]).count(); // 1-based
    let column = simd_utf16_len::utf16_len(&prefix[line_start..]); // 0-based
    let utf16_pos = simd_utf16_len::utf16_len(&prefix[..line_start]) + column;

    let file = span.source().to_string();
    Some((file, line, column, utf16_pos))
  }

  pub fn kind(&self) -> String {
    self.kind.clone()
  }

  /// Code frame of the first label's file, at the `line` and `column` that
  /// [`Self::get_primary_location`] returns.
  pub fn frame(&self, line: usize, column: usize) -> Option<String> {
    let source = self.files.get(self.labels.first()?.span().source())?;
    let frame = render_code_frame(source, line, column);
    (!frame.is_empty()).then_some(frame)
  }

  /// Render many diagnostics at once, sharing per-source work across all of them.
  ///
  /// Rendering diagnostics one by one rebuilds the ariadne `Source` (its rope and
  /// full line index) for the whole file on every diagnostic, so emitting N
  /// diagnostics that point into the same large file is O(N^2) and can appear to
  /// hang the build (#9748). Sharing the `Source` cache and a per-source line table
  /// removes that quadratic factor — the line containing an offset is found in
  /// O(log lines) per diagnostic, then its UTF-16 column is read by scanning only
  /// that one line (see [`ByteLocator::locate_utf16`]).
  ///
  /// One residual cost is outside this batching: ariadne prints the entire labelled
  /// line for every diagnostic, so a pathological input where N diagnostics all land
  /// on one very long line (e.g. a minified bundle) still costs O(N * line_len)
  /// inside ariadne itself — the within-line column scan is bounded by that same
  /// line and is not the deciding factor. The cross-diagnostic O(N^2) blow-up that
  /// actually caused the reported hang is what this removes.
  pub fn render_batch(diagnostics: &[Diagnostic], color: bool) -> Vec<RenderedDiagnostic> {
    // Collect every source referenced by any diagnostic, deduplicated, so the
    // ariadne cache builds each `Source` exactly once.
    let mut files: FxHashMap<DiagnosticFileId, ArcStr> = FxHashMap::default();
    for diagnostic in diagnostics {
      for (id, content) in &diagnostic.files {
        files.entry(id.clone()).or_insert_with(|| content.clone());
      }
    }
    let mut cache = sources(files.iter().map(|(id, content)| (id.clone(), content.clone())));
    let mut locators: FxHashMap<DiagnosticFileId, ByteLocator> = FxHashMap::default();

    diagnostics
      .iter()
      .map(|diagnostic| {
        let primary_location = diagnostic.primary_location_with(&files, &mut locators);
        let message = diagnostic.convert_to_string_with_cache(color, &mut cache);
        RenderedDiagnostic { message, primary_location }
      })
      .collect()
  }

  /// Compute the primary location using a caller-owned per-source line-table
  /// cache, so the line index for each file is built once across many lookups.
  fn primary_location_with(
    &self,
    files: &FxHashMap<DiagnosticFileId, ArcStr>,
    locators: &mut FxHashMap<DiagnosticFileId, ByteLocator>,
  ) -> Option<DiagnosticPrimaryLocation> {
    let first_label = self.labels.first()?;
    let span = first_label.span();
    let source_id = span.source();
    let source = files.get(source_id)?;
    let locator = locators.entry(source_id.clone()).or_insert_with(|| ByteLocator::new(source));
    let (line, column, utf16_position) = locator.locate_utf16(source, span.start());
    Some(DiagnosticPrimaryLocation { line, column, utf16_position })
  }
}

const LINE_TRUNCATE_LENGTH: usize = 120;
const MIN_CHARACTERS_SHOWN_AFTER_LOCATION: usize = 10;
const ELLIPSIS: &str = "...";

fn leading_tabs(value: &str) -> usize {
  value.bytes().take_while(|&byte| byte == b'\t').count()
}

// Only leading tabs are expanded, each to two spaces.
fn tabs_to_spaces(value: &str) -> String {
  let leading_tabs = leading_tabs(value);
  if leading_tabs == 0 {
    return value.to_string();
  }
  let mut expanded = "  ".repeat(leading_tabs);
  expanded.push_str(&value[leading_tabs..]);
  expanded
}

// Mirrors `getCodeFrame` in packages/rolldown/src/utils/code-frame.ts. The column and all
// lengths count UTF-16 code units, like JS strings.
fn render_code_frame(source: &str, line: usize, column: usize) -> String {
  let frame_start = line.saturating_sub(3);
  let mut lines: Vec<&str> =
    source.split('\n').skip(frame_start).take(line + 2 - frame_start).collect();
  let Some(&error_line) = lines.get(line - 1 - frame_start) else {
    return String::new();
  };
  let location_width = column + leading_tabs(error_line).min(column);
  let max_line_length = (location_width + MIN_CHARACTERS_SHOWN_AFTER_LOCATION + ELLIPSIS.len())
    .max(LINE_TRUNCATE_LENGTH);

  while lines.last().is_some_and(|last| last.trim().is_empty()) {
    lines.pop();
  }

  let digits = (frame_start + lines.len()).to_string().len();
  lines
    .iter()
    .enumerate()
    .map(|(index, source_line)| {
      let line_no = frame_start + index + 1;
      let mut displayed = tabs_to_spaces(source_line);
      if simd_utf16_len::utf16_len(&displayed) > max_line_length {
        let mut len = 0;
        let end = displayed
          .char_indices()
          .find(|(_, ch)| {
            len += ch.len_utf16();
            len > max_line_length - ELLIPSIS.len()
          })
          .map_or(displayed.len(), |(index, _)| index);
        displayed.truncate(end);
        displayed.push_str(ELLIPSIS);
      }
      if line_no == line {
        let indent = digits + 2 + location_width;
        format!("{line_no:>digits$}: {displayed}\n{:indent$}^", "")
      } else {
        format!("{line_no:>digits$}: {displayed}")
      }
    })
    .collect::<Vec<_>>()
    .join("\n")
}

/// Primary source location of a diagnostic's first label.
#[derive(Debug, Clone)]
pub struct DiagnosticPrimaryLocation {
  /// 1-based line number.
  pub line: usize,
  /// 0-based UTF-16 column within the line.
  pub column: usize,
  /// UTF-16 code-unit offset from the start of the file.
  pub utf16_position: usize,
}

/// A diagnostic rendered to its display string together with its primary location.
#[derive(Debug, Clone)]
pub struct RenderedDiagnostic {
  pub message: String,
  pub primary_location: Option<DiagnosticPrimaryLocation>,
}

impl Display for Diagnostic {
  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
    self.convert_to_string(false).fmt(f)
  }
}

#[cfg(test)]
mod tests {
  use super::{Diagnostic, Severity, render_code_frame};

  /// Char-by-char reference for [`Diagnostic::get_primary_location`]. Returns
  /// `(line, column, utf16_position)`.
  fn reference_location(source: &str, start: usize) -> (usize, usize, usize) {
    let (mut line, mut column, mut utf16_pos, mut byte_count) = (1, 0, 0, 0);
    for ch in source.chars() {
      if byte_count >= start {
        break;
      }
      if ch == '\n' {
        line += 1;
        column = 0;
      } else {
        column += ch.len_utf16();
      }
      utf16_pos += ch.len_utf16();
      byte_count += ch.len_utf8();
    }
    (line, column, utf16_pos)
  }

  #[test]
  fn get_primary_location_matches_reference() {
    for source in ["", "a\nbc\n\ndef", "café\n😀xy\nz", "\n\n日本語\n💣"] {
      // Every offset, including ones inside a char, plus out-of-range ones.
      #[expect(clippy::cast_possible_truncation)]
      for start in 0..=source.len() as u32 + 2 {
        let mut diagnostic = Diagnostic::new("TEST".into(), String::new(), Severity::Error);
        let file_id = diagnostic.add_file("test.js", source);
        diagnostic.add_label(&file_id, start..start, String::new());
        let (_, line, column, utf16_pos) = diagnostic.get_primary_location().unwrap();
        assert_eq!(
          (line, column, utf16_pos),
          reference_location(source, start as usize),
          "mismatch at start {start} of {source:?}"
        );
      }
    }
  }

  #[test]
  fn code_frame_points_at_location() {
    assert_eq!(render_code_frame("a\nbcd\nef", 2, 1), "1: a\n2: bcd\n    ^\n3: ef");
  }

  #[test]
  fn code_frame_expands_tabs_and_aligns_caret() {
    assert_eq!(render_code_frame("\tfoo", 1, 1), "1:   foo\n     ^");
  }

  #[test]
  fn code_frame_preserves_non_leading_tabs() {
    assert_eq!(render_code_frame("ab\tcd", 1, 3), "1: ab\tcd\n      ^");
  }

  #[test]
  fn code_frame_truncates_long_lines() {
    let frame = render_code_frame(&"x".repeat(200), 1, 0);
    let first_line = frame.lines().next().unwrap();
    assert_eq!(first_line, format!("1: {}...", "x".repeat(117)));
  }

  #[test]
  fn code_frame_drops_trailing_blank_lines() {
    assert_eq!(render_code_frame("a\nb\n\n \n", 2, 0), "1: a\n2: b\n   ^");
  }

  #[test]
  fn frame_returns_none_without_labels() {
    let diagnostic = Diagnostic::new("X".to_string(), "title".to_string(), Severity::Error);
    assert_eq!(diagnostic.frame(1, 0), None);
  }

  #[test]
  fn frame_caret_counts_utf16_units() {
    let mut diagnostic = Diagnostic::new("X".to_string(), "title".to_string(), Severity::Error);
    let file_id = diagnostic.add_file("test.js", "😀x");
    // Byte offset 4 is `x`, which is UTF-16 column 2.
    diagnostic.add_label(&file_id, 4..5, String::new());
    let (_, line, column, _) = diagnostic.get_primary_location().unwrap();
    assert_eq!(diagnostic.frame(line, column).as_deref(), Some("1: 😀x\n     ^"));
  }
}