424 lines
12 KiB
C++
424 lines
12 KiB
C++
#include "codehighlighter.hpp"
|
|
|
|
#include "highlight-queries.hpp"
|
|
|
|
#include <tree_sitter/api.h>
|
|
|
|
#include <QCoreApplication>
|
|
#include <QDir>
|
|
#include <QHash>
|
|
#include <QMap>
|
|
#include <QMutexLocker>
|
|
#include <QPointer>
|
|
#include <QStringList>
|
|
#include <QThreadPool>
|
|
#include <QVariantMap>
|
|
|
|
#include <cstring>
|
|
#include <dlfcn.h>
|
|
|
|
namespace ZShell::llm {
|
|
|
|
namespace hl {
|
|
|
|
enum Role : uint8_t {
|
|
None = 0,
|
|
Comment,
|
|
String,
|
|
StringKey,
|
|
Number,
|
|
Constant,
|
|
Keyword,
|
|
Type,
|
|
Function,
|
|
Method,
|
|
Macro,
|
|
Preproc,
|
|
Operator,
|
|
Property,
|
|
Label,
|
|
Attribute,
|
|
};
|
|
|
|
const char* roleName(Role role) {
|
|
switch (role) {
|
|
case Comment:
|
|
return "comment";
|
|
case String:
|
|
return "string";
|
|
case StringKey:
|
|
return "string.key";
|
|
case Number:
|
|
return "number";
|
|
case Constant:
|
|
return "constant";
|
|
case Keyword:
|
|
return "keyword";
|
|
case Type:
|
|
return "type";
|
|
case Function:
|
|
return "function";
|
|
case Method:
|
|
return "method";
|
|
case Macro:
|
|
return "macro";
|
|
case Preproc:
|
|
return "preproc";
|
|
case Operator:
|
|
return "operator";
|
|
case Property:
|
|
return "property";
|
|
case Label:
|
|
return "label";
|
|
case Attribute:
|
|
return "attribute";
|
|
default:
|
|
return "";
|
|
}
|
|
}
|
|
|
|
using LanguageFn = const TSLanguage* (*)();
|
|
|
|
const QHash<QString, CodeHighlighter::Grammar>& grammars() {
|
|
static const QHash<QString, CodeHighlighter::Grammar> grammars = [] {
|
|
QHash<QString, CodeHighlighter::Grammar> map;
|
|
for (const auto& g : hq::grammars) {
|
|
CodeHighlighter::Grammar grammar;
|
|
for (int i = 0; i < g.nCandidates; ++i) {
|
|
grammar.libs.push_back(g.candidates[i].lib);
|
|
grammar.symbols.push_back(g.candidates[i].symbol);
|
|
}
|
|
for (int i = 0; i < g.nQueries; ++i)
|
|
grammar.queries.push_back(g.queries[i]);
|
|
map.insert(g.id, std::move(grammar));
|
|
}
|
|
return map;
|
|
}();
|
|
return grammars;
|
|
}
|
|
|
|
} // namespace hl
|
|
|
|
CodeHighlighter* CodeHighlighter::s_instance = nullptr;
|
|
|
|
const QHash<QString, QString>& CodeHighlighter::aliases() {
|
|
static const QHash<QString, QString> aliases = [] {
|
|
QHash<QString, QString> map;
|
|
map.insert("c", "c");
|
|
map.insert("h", "c");
|
|
map.insert("cpp", "cpp");
|
|
map.insert("c++", "cpp");
|
|
map.insert("cc", "cpp");
|
|
map.insert("cxx", "cpp");
|
|
map.insert("h++", "cpp");
|
|
map.insert("hpp", "cpp");
|
|
map.insert("hh", "cpp");
|
|
map.insert("python", "python");
|
|
map.insert("py", "python");
|
|
map.insert("javascript", "javascript");
|
|
map.insert("js", "javascript");
|
|
map.insert("jsx", "javascript");
|
|
map.insert("mjs", "javascript");
|
|
map.insert("cjs", "javascript");
|
|
map.insert("typescript", "typescript");
|
|
map.insert("ts", "typescript");
|
|
map.insert("mts", "typescript");
|
|
map.insert("cts", "typescript");
|
|
map.insert("tsx", "tsx");
|
|
map.insert("bash", "bash");
|
|
map.insert("sh", "bash");
|
|
map.insert("shell", "bash");
|
|
map.insert("shellscript", "bash");
|
|
map.insert("shell-session", "bash");
|
|
map.insert("zsh", "bash");
|
|
map.insert("console", "bash");
|
|
map.insert("qml", "qmljs");
|
|
map.insert("qmljs", "qmljs");
|
|
map.insert("json", "json");
|
|
map.insert("jsonc", "json");
|
|
map.insert("rust", "rust");
|
|
map.insert("rs", "rust");
|
|
map.insert("go", "go");
|
|
map.insert("golang", "go");
|
|
map.insert("yaml", "yaml");
|
|
map.insert("yml", "yaml");
|
|
map.insert("toml", "toml");
|
|
map.insert("sql", "sql");
|
|
map.insert("mysql", "sql");
|
|
map.insert("postgres", "sql");
|
|
map.insert("postgresql", "sql");
|
|
map.insert("sqlite", "sql");
|
|
map.insert("sqlite3", "sql");
|
|
return map;
|
|
}();
|
|
return aliases;
|
|
}
|
|
|
|
uint8_t CodeHighlighter::roleFor(const char* name, uint32_t length) {
|
|
const QString n = QString::fromUtf8(name, length);
|
|
if (n == "comment") return hl::Role::Comment;
|
|
if (n.startsWith("string"))
|
|
return n == "string.special.key" ? hl::Role::StringKey
|
|
: hl::Role::String;
|
|
if (n == "escape" || n == "regexp") return hl::Role::String;
|
|
if (n.startsWith("number")) return hl::Role::Number;
|
|
if (n.startsWith("constant") || n == "boolean" || n == "bool" ||
|
|
n.startsWith("character"))
|
|
return hl::Role::Constant;
|
|
if (n.startsWith("keyword")) return hl::Role::Keyword;
|
|
if (n == "type" || n.startsWith("type.")) return hl::Role::Type;
|
|
if (n.startsWith("namespace") || n.startsWith("module") ||
|
|
n == "support.type" || n == "support.namespace")
|
|
return hl::Role::Type;
|
|
if (n.startsWith("function") || n == "constructor" ||
|
|
n.startsWith("support.function"))
|
|
return hl::Role::Function;
|
|
if (n == "method" || n == "method.builtin") return hl::Role::Method;
|
|
if (n.startsWith("macro")) return hl::Role::Macro;
|
|
if (n.startsWith("preproc")) return hl::Role::Preproc;
|
|
if (n == "operator" || n == "punctuation.operator" ||
|
|
n.startsWith("operator.") || n.startsWith("punctuation"))
|
|
return hl::Role::Operator;
|
|
if (n == "property" || n == "field" || n.startsWith("property."))
|
|
return hl::Role::Property;
|
|
if (n == "label") return hl::Role::Label;
|
|
if (n.startsWith("attribute") || n == "annotation")
|
|
return hl::Role::Attribute;
|
|
if (n == "tag" || (n.startsWith("tag.") && n != "tag.delimiter"))
|
|
return hl::Role::Keyword;
|
|
if (n.startsWith("variable")) return hl::Role::Constant;
|
|
if (n.startsWith("support")) return hl::Role::Function;
|
|
return hl::Role::None;
|
|
}
|
|
|
|
const char* CodeHighlighter::roleName(uint8_t role) {
|
|
return hl::roleName(static_cast<hl::Role>(role));
|
|
}
|
|
|
|
QString CodeHighlighter::resolveId(const QString& language) {
|
|
const QString tag = language.trimmed().toLower();
|
|
const QString alias = aliases().value(tag);
|
|
return alias.isEmpty() ? tag : alias; // unknown tags = grammar id
|
|
}
|
|
|
|
QString CodeHighlighter::cacheKey(const QString& id, const QString& code) {
|
|
return id + QLatin1Char('\x01') + QString::number(code.size()) +
|
|
QLatin1Char('\x01') + QString::number(qHash(code));
|
|
}
|
|
|
|
QVariantList CodeHighlighter::lookupSpans(
|
|
const QString& code, const QString& language) const {
|
|
if (code.isEmpty()) return {};
|
|
const QString key = cacheKey(resolveId(language), code);
|
|
QMutexLocker locker(&m_cacheMutex);
|
|
const auto it = m_spanCache.constFind(key);
|
|
if (it == m_spanCache.constEnd() || it->code != code) return {};
|
|
// Most recently used; eviction drops the oldest entries first.
|
|
const qsizetype pos = m_spanCacheOrder.indexOf(key);
|
|
if (pos >= 0) m_spanCacheOrder.move(pos, m_spanCacheOrder.size() - 1);
|
|
return it->spans;
|
|
}
|
|
|
|
void CodeHighlighter::storeSpans(
|
|
const QString& code,
|
|
const QString& language,
|
|
const QVariantList& spans) const {
|
|
if (spans.isEmpty() || code.isEmpty()) return;
|
|
static constexpr int kMaxEntries = 32;
|
|
static constexpr int kMaxBytes = 1024 * 1024;
|
|
const QString key = cacheKey(resolveId(language), code);
|
|
const int bytes = static_cast<int>(code.toUtf8().size());
|
|
QMutexLocker locker(&m_cacheMutex);
|
|
auto it = m_spanCache.find(key);
|
|
if (it != m_spanCache.end()) {
|
|
m_spanCacheBytes -= static_cast<int>(it->code.toUtf8().size());
|
|
m_spanCache.erase(it);
|
|
m_spanCacheOrder.removeAll(key);
|
|
}
|
|
while (m_spanCacheOrder.size() >= kMaxEntries ||
|
|
m_spanCacheBytes + bytes > kMaxBytes) {
|
|
if (m_spanCacheOrder.isEmpty()) break;
|
|
const QString oldest = m_spanCacheOrder.takeFirst();
|
|
m_spanCacheBytes -=
|
|
static_cast<int>(m_spanCache.value(oldest).code.toUtf8().size());
|
|
m_spanCache.remove(oldest);
|
|
}
|
|
m_spanCache.insert(key, SpanCacheEntry{code, spans});
|
|
m_spanCacheOrder.append(key);
|
|
m_spanCacheBytes += bytes;
|
|
}
|
|
|
|
void CodeHighlighter::highlight(
|
|
const QString& code, const QString& language, QObject* target, int token) {
|
|
QPointer<QObject> targetGuard(target);
|
|
const QVariantList cached = lookupSpans(code, language);
|
|
if (!cached.isEmpty()) {
|
|
QMetaObject::invokeMethod(
|
|
targetGuard,
|
|
"onHighlightSpans",
|
|
Qt::DirectConnection,
|
|
Q_ARG(QVariant, token),
|
|
Q_ARG(QVariant, cached));
|
|
return;
|
|
}
|
|
QThreadPool::globalInstance()->start(
|
|
[this, target, token, code, language]() {
|
|
const QVariantList spans = doHighlight(code, language);
|
|
storeSpans(code, language, spans);
|
|
QPointer<QObject> guard(target);
|
|
QMetaObject::invokeMethod(
|
|
QCoreApplication::instance(),
|
|
[guard, token, spans]() {
|
|
if (!guard) return;
|
|
QMetaObject::invokeMethod(
|
|
guard,
|
|
"onHighlightSpans",
|
|
Q_ARG(QVariant, token),
|
|
Q_ARG(QVariant, spans));
|
|
},
|
|
Qt::QueuedConnection);
|
|
});
|
|
}
|
|
|
|
QVariantList CodeHighlighter::doHighlight(
|
|
const QString& code, const QString& language) const {
|
|
QVariantList spans;
|
|
if (code.isEmpty()) return spans;
|
|
|
|
const QString id = resolveId(language);
|
|
const Grammar& grammar = hl::grammars().value(id);
|
|
if (grammar.libs.empty()) return spans;
|
|
|
|
static constexpr size_t kMaxBytes = 512 * 1024;
|
|
const QByteArray utf8 = code.toUtf8();
|
|
if (static_cast<size_t>(utf8.size()) > kMaxBytes) return spans;
|
|
|
|
const TSLanguage* lang = nullptr;
|
|
TSQuery* query = nullptr;
|
|
{
|
|
QMutexLocker locker(&m_stateMutex);
|
|
auto& state = m_states[id];
|
|
if (!state || (!state->lang && !state->bad)) {
|
|
std::shared_ptr<State> fresh = std::make_shared<State>();
|
|
bool abiMismatch = false;
|
|
for (size_t i = 0; i < grammar.libs.size(); ++i) {
|
|
void* lib =
|
|
dlopen(grammar.libs[i].c_str(), RTLD_NOW | RTLD_LOCAL);
|
|
if (!lib) continue;
|
|
auto* symbol = reinterpret_cast<hl::LanguageFn>(
|
|
dlsym(lib, grammar.symbols[i].c_str()));
|
|
if (!symbol) {
|
|
dlclose(lib);
|
|
continue;
|
|
}
|
|
const TSLanguage* candidate = symbol();
|
|
const uint32_t version = ts_language_abi_version(candidate);
|
|
if (version < TREE_SITTER_MIN_COMPATIBLE_LANGUAGE_VERSION ||
|
|
version > TREE_SITTER_LANGUAGE_VERSION) {
|
|
dlclose(lib);
|
|
abiMismatch = true;
|
|
continue;
|
|
}
|
|
fresh->lib = lib;
|
|
fresh->lang = candidate;
|
|
break;
|
|
}
|
|
if (fresh->lang) {
|
|
for (const char* source : grammar.queries) {
|
|
TSQueryError errorType = TSQueryErrorNone;
|
|
uint32_t errorOffset = 0;
|
|
TSQuery* candidate = ts_query_new(
|
|
static_cast<const TSLanguage*>(fresh->lang),
|
|
source,
|
|
static_cast<uint32_t>(std::strlen(source)),
|
|
&errorOffset,
|
|
&errorType);
|
|
if (!candidate) continue;
|
|
fresh->query = candidate;
|
|
break;
|
|
}
|
|
if (!fresh->query) fresh->bad = true;
|
|
} else if (abiMismatch) {
|
|
fresh->bad = true;
|
|
}
|
|
if (fresh->lang || fresh->bad) state = std::move(fresh);
|
|
}
|
|
if (!state || state->bad || !state->lang) return spans;
|
|
lang = static_cast<const TSLanguage*>(state->lang);
|
|
query = static_cast<TSQuery*>(state->query);
|
|
}
|
|
|
|
TSParser* parser = ts_parser_new();
|
|
ts_parser_set_language(parser, lang);
|
|
TSTree* tree = ts_parser_parse_string(
|
|
parser, nullptr, utf8.constData(), static_cast<uint32_t>(utf8.size()));
|
|
if (!tree) {
|
|
ts_parser_delete(parser);
|
|
return spans;
|
|
}
|
|
|
|
TSQueryCursor* cursor = ts_query_cursor_new();
|
|
ts_query_cursor_exec(cursor, query, ts_tree_root_node(tree));
|
|
|
|
const uint32_t size = static_cast<uint32_t>(utf8.size());
|
|
std::vector<uint8_t> kinds(size, 0);
|
|
|
|
std::vector<uint32_t> cu(size + 1, 0);
|
|
for (uint32_t b = 0; b < size; ++b) {
|
|
cu[b + 1] = cu[b];
|
|
const unsigned char c = static_cast<unsigned char>(utf8[b]);
|
|
if (c < 0x80)
|
|
cu[b + 1] += 1;
|
|
else if (c < 0xC0)
|
|
; // continuation byte
|
|
else if (c < 0xF0)
|
|
cu[b + 1] += 1;
|
|
else
|
|
cu[b + 1] += 2;
|
|
}
|
|
|
|
TSQueryMatch match;
|
|
uint32_t captureIndex = 0;
|
|
while (ts_query_cursor_next_capture(cursor, &match, &captureIndex)) {
|
|
const TSQueryCapture& capture = match.captures[captureIndex];
|
|
uint32_t nameLength = 0;
|
|
const char* name =
|
|
ts_query_capture_name_for_id(query, capture.index, &nameLength);
|
|
const uint8_t role = roleFor(name, nameLength);
|
|
if (role == 0) continue;
|
|
const uint32_t start = ts_node_start_byte(capture.node);
|
|
const uint32_t end = ts_node_end_byte(capture.node);
|
|
if (end <= start || end > size) continue;
|
|
std::fill(kinds.begin() + start, kinds.begin() + end, role);
|
|
}
|
|
|
|
uint32_t position = 0;
|
|
while (position < size) {
|
|
if (kinds[position] == 0) {
|
|
++position;
|
|
continue;
|
|
}
|
|
const uint8_t role = kinds[position];
|
|
const uint32_t start = position;
|
|
while (position < size && kinds[position] == role)
|
|
++position;
|
|
QVariantMap span;
|
|
span.insert("start", static_cast<int>(cu[start]));
|
|
span.insert("length", static_cast<int>(cu[position] - cu[start]));
|
|
span.insert("kind", roleName(role));
|
|
spans.append(span);
|
|
}
|
|
|
|
ts_query_cursor_delete(cursor);
|
|
ts_tree_delete(tree);
|
|
ts_parser_delete(parser);
|
|
return spans;
|
|
}
|
|
|
|
CodeHighlighter* CodeHighlighter::create(QQmlEngine*, QJSEngine*) {
|
|
if (!s_instance) s_instance = new CodeHighlighter();
|
|
return s_instance;
|
|
}
|
|
|
|
} // namespace ZShell::llm
|