add markdown parsing + latex + tree-sitter highlighting for codeblocks in llm responses
This commit is contained in:
@@ -0,0 +1,384 @@
|
||||
#include "codehighlighter.hpp"
|
||||
|
||||
#include "highlight-queries.hpp"
|
||||
|
||||
#include <tree_sitter/api.h>
|
||||
|
||||
#include <QHash>
|
||||
#include <QMap>
|
||||
#include <QVariantMap>
|
||||
|
||||
#include <dlfcn.h>
|
||||
|
||||
namespace ZShell::llm {
|
||||
|
||||
namespace hl {
|
||||
|
||||
// Role ids; 0 means "no color".
|
||||
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* (*)();
|
||||
|
||||
// Query sources, indexed by Grammar::queries.
|
||||
struct QuerySources {
|
||||
std::vector<std::string> sources;
|
||||
int c, cpp, python, javascript, typescriptExtended, typescript, bash, json, rust, go, yaml, toml, sql, cppExtended;
|
||||
};
|
||||
|
||||
const QuerySources& querySources() {
|
||||
static const QuerySources sources = [] {
|
||||
QuerySources qs;
|
||||
// javascript + typescript concatenated: the TS grammar reuses the
|
||||
// JS node names, so the JS query usually compiles against it and
|
||||
// gives full coverage; the TS-only query is the fallback.
|
||||
const std::string tsExtended =
|
||||
std::string(hq::javascript) + "\n" + std::string(hq::typescript);
|
||||
// Same for cpp: the C++ grammar is a superset of C, and the C++
|
||||
// query only covers the C++-specific delta. Base C coverage
|
||||
// (keywords, types, calls, strings, comments) comes from the C
|
||||
// query.
|
||||
const std::string cppExtended =
|
||||
std::string(hq::c) + "\n" + std::string(hq::cpp);
|
||||
qs.sources.reserve(14);
|
||||
qs.sources.push_back(hq::c);
|
||||
qs.sources.push_back(hq::cpp);
|
||||
qs.sources.push_back(hq::python);
|
||||
qs.sources.push_back(hq::javascript);
|
||||
qs.sources.push_back(tsExtended);
|
||||
qs.sources.push_back(hq::typescript);
|
||||
qs.sources.push_back(hq::bash);
|
||||
qs.sources.push_back(hq::json);
|
||||
qs.sources.push_back(hq::rust);
|
||||
qs.sources.push_back(hq::go);
|
||||
qs.sources.push_back(hq::yaml);
|
||||
qs.sources.push_back(hq::toml);
|
||||
qs.sources.push_back(hq::sql);
|
||||
qs.sources.push_back(cppExtended);
|
||||
qs.c = 0;
|
||||
qs.cpp = 1;
|
||||
qs.python = 2;
|
||||
qs.javascript = 3;
|
||||
qs.typescriptExtended = 4;
|
||||
qs.typescript = 5;
|
||||
qs.bash = 6;
|
||||
qs.json = 7;
|
||||
qs.rust = 8;
|
||||
qs.go = 9;
|
||||
qs.yaml = 10;
|
||||
qs.toml = 11;
|
||||
qs.sql = 12;
|
||||
qs.cppExtended = 13;
|
||||
return qs;
|
||||
}();
|
||||
return sources;
|
||||
}
|
||||
|
||||
const QHash<QString, CodeHighlighter::Grammar>& grammars() {
|
||||
static const QHash<QString, CodeHighlighter::Grammar> grammars = [] {
|
||||
const auto& qs = querySources();
|
||||
QHash<QString, CodeHighlighter::Grammar> map;
|
||||
map.insert("c", {"libtree-sitter-c.so", "tree_sitter_c", {qs.c}});
|
||||
map.insert("cpp", {"libtree-sitter-cpp.so", "tree_sitter_cpp", {qs.cppExtended, qs.cpp}});
|
||||
map.insert("python", {"libtree-sitter-python.so", "tree_sitter_python", {qs.python}});
|
||||
map.insert("javascript", {"libtree-sitter-javascript.so", "tree_sitter_javascript", {qs.javascript}});
|
||||
map.insert("typescript", {"libtree-sitter-typescript.so", "tree_sitter_typescript", {qs.typescriptExtended, qs.typescript}});
|
||||
map.insert("tsx", {"libtree-sitter-tsx.so", "tree_sitter_tsx", {qs.typescriptExtended, qs.typescript}});
|
||||
map.insert("bash", {"libtree-sitter-bash.so", "tree_sitter_bash", {qs.bash}});
|
||||
map.insert("json", {"libtree-sitter-json.so", "tree_sitter_json", {qs.json}});
|
||||
map.insert("rust", {"libtree-sitter-rust.so", "tree_sitter_rust", {qs.rust}});
|
||||
map.insert("go", {"libtree-sitter-go.so", "tree_sitter_go", {qs.go}});
|
||||
map.insert("yaml", {"libtree-sitter-yaml.so", "tree_sitter_yaml", {qs.yaml}});
|
||||
map.insert("toml", {"libtree-sitter-toml.so", "tree_sitter_toml", {qs.toml}});
|
||||
map.insert("sql", {"libtree-sitter-sql.so", "tree_sitter_sql", {qs.sql}});
|
||||
return map;
|
||||
}();
|
||||
return grammars;
|
||||
}
|
||||
|
||||
} // namespace hl
|
||||
|
||||
CodeHighlighter* CodeHighlighter::s_instance = nullptr;
|
||||
|
||||
const QHash<QString, QString>& CodeHighlighter::aliases() {
|
||||
// Language tags as written in code fences (and common variants) to
|
||||
// grammar id.
|
||||
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("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;
|
||||
}
|
||||
|
||||
const std::vector<std::string>& CodeHighlighter::querySources() const {
|
||||
return hl::querySources().sources;
|
||||
}
|
||||
|
||||
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"))
|
||||
return hl::Role::Constant;
|
||||
if (n.startsWith("keyword"))
|
||||
return hl::Role::Keyword;
|
||||
if (n == "type" || n.startsWith("type."))
|
||||
return hl::Role::Type;
|
||||
if (n == "namespace" || n == "module")
|
||||
return hl::Role::Type;
|
||||
if (n.startsWith("function") || n == "constructor")
|
||||
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."))
|
||||
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;
|
||||
return hl::Role::None;
|
||||
}
|
||||
|
||||
const char* CodeHighlighter::roleName(uint8_t role) {
|
||||
return hl::roleName(static_cast<hl::Role>(role));
|
||||
}
|
||||
|
||||
QVariantList CodeHighlighter::highlight(const QString& code, const QString& language) const {
|
||||
QVariantList spans;
|
||||
if (code.isEmpty())
|
||||
return spans;
|
||||
|
||||
const QString id = aliases().value(language.trimmed().toLower());
|
||||
if (id.isEmpty())
|
||||
return spans;
|
||||
const Grammar& grammar = hl::grammars().value(id);
|
||||
|
||||
// Guard against pathological blocks; highlighting is best-effort.
|
||||
static constexpr size_t kMaxBytes = 512 * 1024;
|
||||
const QByteArray utf8 = code.toUtf8();
|
||||
if (static_cast<size_t>(utf8.size()) > kMaxBytes)
|
||||
return spans;
|
||||
|
||||
State& state = m_states[id];
|
||||
if (state.bad)
|
||||
return spans;
|
||||
|
||||
if (!state.lang) {
|
||||
// Missing library is retriable (it may be installed while the
|
||||
// shell runs); ABI/query failures below are not.
|
||||
state.lib = dlopen(grammar.lib.toUtf8().constData(), RTLD_NOW | RTLD_LOCAL);
|
||||
if (!state.lib)
|
||||
return spans;
|
||||
auto* symbol = reinterpret_cast<hl::LanguageFn>(
|
||||
dlsym(state.lib, grammar.symbol.toUtf8().constData()));
|
||||
if (!symbol) {
|
||||
dlclose(state.lib);
|
||||
state.lib = nullptr;
|
||||
return spans;
|
||||
}
|
||||
const TSLanguage* lang = symbol();
|
||||
const uint32_t version = ts_language_abi_version(lang);
|
||||
if (version < TREE_SITTER_MIN_COMPATIBLE_LANGUAGE_VERSION ||
|
||||
version > TREE_SITTER_LANGUAGE_VERSION) {
|
||||
state.bad = true;
|
||||
return spans;
|
||||
}
|
||||
state.lang = lang;
|
||||
}
|
||||
|
||||
if (!state.query) {
|
||||
// Candidates in priority order; first that compiles wins.
|
||||
for (const int candidate : grammar.queries) {
|
||||
const std::string& source =
|
||||
querySources()[static_cast<size_t>(candidate)];
|
||||
TSQueryError errorType = TSQueryErrorNone;
|
||||
uint32_t errorOffset = 0;
|
||||
TSQuery* query = ts_query_new(
|
||||
static_cast<const TSLanguage*>(state.lang),
|
||||
source.data(),
|
||||
static_cast<uint32_t>(source.size()),
|
||||
&errorOffset,
|
||||
&errorType);
|
||||
if (!query)
|
||||
continue;
|
||||
state.query = query;
|
||||
break;
|
||||
}
|
||||
if (!state.query) {
|
||||
state.bad = true;
|
||||
return spans;
|
||||
}
|
||||
}
|
||||
|
||||
TSParser* parser = ts_parser_new();
|
||||
ts_parser_set_language(parser, static_cast<const TSLanguage*>(state.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;
|
||||
}
|
||||
|
||||
TSQuery* query = static_cast<TSQuery*>(state.query);
|
||||
TSQueryCursor* cursor = ts_query_cursor_new();
|
||||
ts_query_cursor_exec(cursor, query, ts_tree_root_node(tree));
|
||||
|
||||
// Per-byte winner table: captures arrive in document order and later
|
||||
// captures overwrite earlier ones (tree-sitter highlight convention).
|
||||
const uint32_t size = static_cast<uint32_t>(utf8.size());
|
||||
std::vector<uint8_t> kinds(size, 0);
|
||||
|
||||
// QML slices the code by UTF-16 code unit, so spans must be in code
|
||||
// units, not bytes. cu[b] = code units before byte b.
|
||||
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; // 2/3-byte lead -> BMP -> one unit
|
||||
else
|
||||
cu[b + 1] += 2; // 4-byte lead -> surrogate pair
|
||||
}
|
||||
|
||||
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
|
||||
Reference in New Issue
Block a user