z-bar-qt/Plugins/ZShell/Llm/codehighlighter.cpp

#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 {

// 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* (*)();

// Grammar registry generated by CMake from the installed grammars and
// their highlight queries (see CMakeLists.txt).
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() {
	// 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("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;
	// HTML tag names, CSS variables and friends.
	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));
}

void CodeHighlighter::highlight(
	const QString& code, const QString& language, QObject* target, int token) {
	QThreadPool::globalInstance()->start([this, target, token, code, language]() {
		const QVariantList spans = doHighlight(code, language);
		// The target item may be long gone by now (delegates are
		// recreated constantly while chats load); a destroyed target is
		// simply skipped. Deliver through the app instance (never
		// destroyed) and re-check there: posting to `target` from the
		// pool thread would race with its destruction.
		QPointer<QObject> guard(target);
		QMetaObject::invokeMethod(
			QCoreApplication::instance(),
			[guard, token, spans]() {
				if (!guard)
					return;
			// QML functions are only invokable by their generic
			// QVariant overload, so pass untyped arguments.
			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 tag = language.trimmed().toLower();
	const QString id = [&] {
		const QString alias = aliases().value(tag);
		return alias.isEmpty() ? tag : alias; // unknown tags = grammar id
	}();
	const Grammar& grammar = hl::grammars().value(id);
	if (grammar.libs.empty())
		return spans;

	// 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;

	const TSLanguage* lang = nullptr;
	TSQuery* query = nullptr;
	{
		QMutexLocker locker(&m_stateMutex);
		auto& state = m_states[id];
		// A missing library is retriable (it may be installed while the
		// shell runs); an ABI mismatch on every candidate is not. Cache
		// successes and permanent failures; leave retriable misses out.
		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) {
				// Candidates in priority order; first that compiles wins.
				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));

	// 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