TheRock/compiler/amd-llvm/lldb/source/Plugins/Highlighter/TreeSitter/TreeSitterHighlighter.cpp

//===----------------------------------------------------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#include "TreeSitterHighlighter.h"
#include "lldb/Utility/LLDBLog.h"
#include "lldb/Utility/Log.h"
#include "lldb/Utility/StreamString.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/StringSwitch.h"

using namespace lldb_private;

TreeSitterHighlighter::TSState::~TSState() {
  if (query)
    ts_query_delete(query);
  if (parser)
    ts_parser_delete(parser);
}

TreeSitterHighlighter::TSState::operator bool() const {
  return parser && query;
}

TreeSitterHighlighter::TSState &TreeSitterHighlighter::GetTSState() const {
  if (m_ts_state)
    return *m_ts_state;

  Log *log = GetLog(LLDBLog::Source);

  m_ts_state.emplace();
  m_ts_state->parser = ts_parser_new();
  if (!m_ts_state->parser) {
    LLDB_LOG(log, "Creating tree-sitter parser failed for {0}", GetName());
    return *m_ts_state;
  }

  const TSLanguage *language = GetLanguage();
  if (!language || !ts_parser_set_language(m_ts_state->parser, language)) {
    LLDB_LOG(log, "Creating tree-sitter language failed for {0}", GetName());
    return *m_ts_state;
  }

  llvm::StringRef query_source = GetHighlightQuery();
  uint32_t error_offset = 0;
  TSQueryError error_type = TSQueryErrorNone;
  m_ts_state->query = ts_query_new(language, query_source.data(),
                                   static_cast<uint32_t>(query_source.size()),
                                   &error_offset, &error_type);
  if (!m_ts_state->query || error_type != TSQueryErrorNone) {
    LLDB_LOG(log,
             "Creating tree-sitter query failed for {0} with error {1}: {2}",
             GetName(), error_type, query_source.substr(error_offset, 64));
    // If we have an error but a valid query, we need to reset the object to
    // (1) avoid it looking valid and (2) release the parser.
    m_ts_state.emplace();
  }

  return *m_ts_state;
}

HighlightStyle::ColorStyle
TreeSitterHighlighter::GetStyleForCapture(llvm::StringRef capture_name,
                                          const HighlightStyle &options) const {
  return llvm::StringSwitch<HighlightStyle::ColorStyle>(capture_name)
      .Case("comment", options.comment)
      .Case("keyword", options.keyword)
      .Case("operator", options.operators)
      .Case("type", options.keyword)
      .Case("punctuation.delimiter.comma", options.comma)
      .Case("punctuation.delimiter.colon", options.colon)
      .Case("punctuation.delimiter.semicolon", options.semicolons)
      .Case("punctuation.bracket.square", options.square_brackets)
      .Cases({"keyword.directive", "preproc"}, options.pp_directive)
      .Cases({"string", "string.literal"}, options.string_literal)
      .Cases({"number", "number.literal", "constant.numeric"},
             options.scalar_literal)
      .Cases({"identifier", "variable", "function"}, options.identifier)
      .Cases({"punctuation.bracket.curly", "punctuation.brace"}, options.braces)
      .Cases({"punctuation.bracket.round", "punctuation.bracket",
              "punctuation.paren"},
             options.parentheses)
      .Default({});
}

void TreeSitterHighlighter::HighlightRange(
    const HighlightStyle &options, llvm::StringRef text, uint32_t start_byte,
    uint32_t end_byte, const HighlightStyle::ColorStyle &style,
    std::optional<size_t> cursor_pos, bool &highlighted_cursor,
    Stream &s) const {

  if (start_byte >= end_byte || start_byte >= text.size())
    return;

  end_byte = std::min(end_byte, static_cast<uint32_t>(text.size()));

  llvm::StringRef range = text.substr(start_byte, end_byte - start_byte);

  auto print = [&](llvm::StringRef str) {
    if (style)
      style.Apply(s, str);
    else
      s << str;
  };

  // Check if cursor is within this range.
  if (cursor_pos && *cursor_pos >= start_byte && *cursor_pos < end_byte &&
      !highlighted_cursor) {
    highlighted_cursor = true;

    // Split range around cursor position.
    const size_t cursor_in_range = *cursor_pos - start_byte;

    // Print everything before the cursor.
    if (cursor_in_range > 0) {
      llvm::StringRef before = range.substr(0, cursor_in_range);
      print(before);
    }

    // Print the cursor itself.
    if (cursor_in_range < range.size()) {
      StreamString cursor_str;
      llvm::StringRef cursor_char = range.substr(cursor_in_range, 1);
      if (style)
        style.Apply(cursor_str, cursor_char);
      else
        cursor_str << cursor_char;
      options.selected.Apply(s, cursor_str.GetString());
    }

    // Print everything after the cursor.
    if (cursor_in_range + 1 < range.size()) {
      llvm::StringRef after = range.substr(cursor_in_range + 1);
      print(after);
    }
  } else {
    // No cursor in this range, apply style directly.
    print(range);
  }
}

void TreeSitterHighlighter::Highlight(const HighlightStyle &options,
                                      llvm::StringRef line,
                                      std::optional<size_t> cursor_pos,
                                      llvm::StringRef previous_lines,
                                      Stream &s) const {
  auto unformatted = [&]() -> void { s << line; };

  TSState &ts_state = GetTSState();
  if (!ts_state)
    return unformatted();

  std::string source = previous_lines.str() + line.str();
  TSTree *tree =
      ts_parser_parse_string(ts_state.parser, nullptr, source.c_str(),
                             static_cast<uint32_t>(source.size()));
  if (!tree)
    return unformatted();

  TSQueryCursor *cursor = ts_query_cursor_new();
  assert(cursor);

  llvm::scope_exit delete_cusor([&] { ts_query_cursor_delete(cursor); });

  TSNode root_node = ts_tree_root_node(tree);
  ts_query_cursor_exec(cursor, ts_state.query, root_node);

  // Collect all matches and their byte ranges.
  std::vector<HLRange> highlights;
  TSQueryMatch match;
  uint32_t capture_index;
  while (ts_query_cursor_next_capture(cursor, &match, &capture_index)) {
    TSQueryCapture capture = match.captures[capture_index];

    uint32_t capture_name_len = 0;
    const char *capture_name = ts_query_capture_name_for_id(
        ts_state.query, capture.index, &capture_name_len);

    HighlightStyle::ColorStyle style = GetStyleForCapture(
        llvm::StringRef(capture_name, capture_name_len), options);

    TSNode node = capture.node;
    uint32_t start = ts_node_start_byte(node);
    uint32_t end = ts_node_end_byte(node);

    if (style && start < end)
      highlights.push_back({start, end, style});
  }

  std::sort(highlights.begin(), highlights.end(),
            [](const HLRange &a, const HLRange &b) {
              if (a.start_byte != b.start_byte)
                return a.start_byte < b.start_byte;
              // Prefer shorter matches. For example, if we have an expression
              // consisting of a variable and a property, we want to highlight
              // them as individual components.
              return (b.end_byte - b.start_byte) > (a.end_byte - a.start_byte);
            });

  uint32_t current_pos = 0;
  bool highlighted_cursor = false;

  for (const auto &h : highlights) {
    // Skip over highlights that start before our current position, which means
    // there's overlap.
    if (h.start_byte < current_pos)
      continue;

    // Output any unhighlighted text before this highlight.
    if (current_pos < h.start_byte) {
      HighlightRange(options, line, current_pos, h.start_byte, {}, cursor_pos,
                     highlighted_cursor, s);
      current_pos = h.start_byte;
    }

    // Output the highlighted range.
    HighlightRange(options, line, h.start_byte, h.end_byte, h.style, cursor_pos,
                   highlighted_cursor, s);
    current_pos = h.end_byte;
  }

  // Output any remaining unhighlighted text.
  if (current_pos < line.size()) {
    HighlightRange(options, line, current_pos,
                   static_cast<uint32_t>(line.size()), {}, cursor_pos,
                   highlighted_cursor, s);
  }
}