should be it

external/duckdb/tools/shell/linenoise/highlighting.cpp

@@ -0,0 +1,270 @@
+#include "linenoise.hpp"
+#include "linenoise.h"
+#include "highlighting.hpp"
+#include "duckdb/parser/parser.hpp"
+#include "duckdb/common/string.hpp"
+
+#if defined(_WIN32) || defined(__WIN32__) || defined(WIN32)
+// disable highlighting on windows (for now?)
+#define DISABLE_HIGHLIGHT
+#endif
+
+namespace duckdb {
+
+#ifdef DISABLE_HIGHLIGHT
+static int enableHighlighting = 0;
+#else
+static int enableHighlighting = 1;
+#endif
+struct Color {
+	const char *color_name;
+	const char *highlight;
+};
+static Color terminal_colors[] = {{"red", "\033[31m"},           {"green", "\033[32m"},
+                                  {"yellow", "\033[33m"},        {"blue", "\033[34m"},
+                                  {"magenta", "\033[35m"},       {"cyan", "\033[36m"},
+                                  {"white", "\033[37m"},         {"brightblack", "\033[90m"},
+                                  {"brightred", "\033[91m"},     {"brightgreen", "\033[92m"},
+                                  {"brightyellow", "\033[93m"},  {"brightblue", "\033[94m"},
+                                  {"brightmagenta", "\033[95m"}, {"brightcyan", "\033[96m"},
+                                  {"brightwhite", "\033[97m"},   {nullptr, nullptr}};
+static std::string bold = "\033[1m";
+static std::string underline = "\033[4m";
+static std::string keyword = "\033[32m";
+static std::string continuation_selected = "\033[32m";
+static std::string constant = "\033[33m";
+static std::string continuation = "\033[90m";
+static std::string comment = "\033[90m";
+static std::string error = "\033[31m";
+static std::string reset = "\033[00m";
+
+void Highlighting::Enable() {
+	enableHighlighting = 1;
+}
+
+void Highlighting::Disable() {
+	enableHighlighting = 0;
+}
+
+bool Highlighting::IsEnabled() {
+	return enableHighlighting;
+}
+
+const char *Highlighting::GetColorOption(const char *option) {
+	size_t index = 0;
+	while (terminal_colors[index].color_name) {
+		if (strcmp(terminal_colors[index].color_name, option) == 0) {
+			return terminal_colors[index].highlight;
+		}
+		index++;
+	}
+	return nullptr;
+}
+
+void Highlighting::SetHighlightingColor(HighlightingType type, const char *color) {
+	switch (type) {
+	case HighlightingType::KEYWORD:
+		keyword = color;
+		break;
+	case HighlightingType::CONSTANT:
+		constant = color;
+		break;
+	case HighlightingType::COMMENT:
+		comment = color;
+		break;
+	case HighlightingType::ERROR:
+		error = color;
+		break;
+	case HighlightingType::CONTINUATION:
+		continuation = color;
+		break;
+	case HighlightingType::CONTINUATION_SELECTED:
+		continuation_selected = color;
+		break;
+	}
+}
+
+static tokenType convertToken(duckdb::SimplifiedTokenType token_type) {
+	switch (token_type) {
+	case duckdb::SimplifiedTokenType::SIMPLIFIED_TOKEN_IDENTIFIER:
+		return tokenType::TOKEN_IDENTIFIER;
+	case duckdb::SimplifiedTokenType::SIMPLIFIED_TOKEN_NUMERIC_CONSTANT:
+		return tokenType::TOKEN_NUMERIC_CONSTANT;
+	case duckdb::SimplifiedTokenType::SIMPLIFIED_TOKEN_STRING_CONSTANT:
+		return tokenType::TOKEN_STRING_CONSTANT;
+	case duckdb::SimplifiedTokenType::SIMPLIFIED_TOKEN_OPERATOR:
+		return tokenType::TOKEN_OPERATOR;
+	case duckdb::SimplifiedTokenType::SIMPLIFIED_TOKEN_KEYWORD:
+		return tokenType::TOKEN_KEYWORD;
+	case duckdb::SimplifiedTokenType::SIMPLIFIED_TOKEN_COMMENT:
+		return tokenType::TOKEN_COMMENT;
+	default:
+		throw duckdb::InternalException("Unrecognized token type");
+	}
+}
+
+static vector<highlightToken> GetParseTokens(char *buf, size_t len) {
+	string sql(buf, len);
+	auto parseTokens = duckdb::Parser::Tokenize(sql);
+
+	vector<highlightToken> tokens;
+	for (auto &token : parseTokens) {
+		highlightToken new_token;
+		new_token.type = convertToken(token.type);
+		new_token.start = token.start;
+		tokens.push_back(new_token);
+	}
+
+	if (!tokens.empty() && tokens[0].start > 0) {
+		highlightToken new_token;
+		new_token.type = tokenType::TOKEN_IDENTIFIER;
+		new_token.start = 0;
+		tokens.insert(tokens.begin(), new_token);
+	}
+	if (tokens.empty() && sql.size() > 0) {
+		highlightToken new_token;
+		new_token.type = tokenType::TOKEN_IDENTIFIER;
+		new_token.start = 0;
+		tokens.push_back(new_token);
+	}
+	return tokens;
+}
+
+static vector<highlightToken> GetDotCommandTokens(char *buf, size_t len) {
+	vector<highlightToken> tokens;
+
+	// identifier token for the dot command itself
+	highlightToken dot_token;
+	dot_token.type = tokenType::TOKEN_KEYWORD;
+	dot_token.start = 0;
+	tokens.push_back(dot_token);
+
+	for (idx_t i = 0; i + 1 < len; i++) {
+		if (Linenoise::IsSpace(buf[i])) {
+			highlightToken argument_token;
+			argument_token.type = tokenType::TOKEN_STRING_CONSTANT;
+			argument_token.start = i + 1;
+			tokens.push_back(argument_token);
+		}
+	}
+	return tokens;
+}
+
+vector<highlightToken> Highlighting::Tokenize(char *buf, size_t len, bool is_dot_command, searchMatch *match) {
+	vector<highlightToken> tokens;
+	if (!is_dot_command) {
+		// SQL query - use parser to obtain tokens
+		tokens = GetParseTokens(buf, len);
+	} else {
+		// . command
+		tokens = GetDotCommandTokens(buf, len);
+	}
+	if (match) {
+		// we have a search match - insert it into the token list
+		// we want to insert a search token with start = match_start, end = match_end
+		// first figure out which token type we would have at match_end (if any)
+		for (size_t i = 0; i + 1 < tokens.size(); i++) {
+			if (tokens[i].start <= match->match_start && tokens[i + 1].start >= match->match_start) {
+				// this token begins after the search position, insert the token here
+				size_t token_position = i + 1;
+				auto end_type = tokens[i].type;
+				if (tokens[i].start == match->match_start) {
+					// exact start: only set the search match
+					tokens[i].search_match = true;
+				} else {
+					// non-exact start: add a new token
+					highlightToken search_token;
+					search_token.type = tokens[i].type;
+					search_token.start = match->match_start;
+					search_token.search_match = true;
+					tokens.insert(tokens.begin() + token_position, search_token);
+					token_position++;
+				}
+
+				// move forwards
+				while (token_position < tokens.size() && tokens[token_position].start < match->match_end) {
+					// this token is
+					// mark this token as a search token
+					end_type = tokens[token_position].type;
+					tokens[token_position].search_match = true;
+					token_position++;
+				}
+				if (token_position >= tokens.size() || tokens[token_position].start > match->match_end) {
+					// insert the token that marks the end of the search
+					highlightToken end_token;
+					end_token.type = end_type;
+					end_token.start = match->match_end;
+					tokens.insert(tokens.begin() + token_position, end_token);
+					token_position++;
+				}
+				break;
+			}
+		}
+	}
+	return tokens;
+}
+
+string Highlighting::HighlightText(char *buf, size_t len, size_t start_pos, size_t end_pos,
+                                   const vector<highlightToken> &tokens) {
+	duckdb::stringstream ss;
+	size_t prev_pos = 0;
+	for (size_t i = 0; i < tokens.size(); i++) {
+		size_t next = i + 1 < tokens.size() ? tokens[i + 1].start : len;
+		if (next < start_pos) {
+			// this token is not rendered at all
+			continue;
+		}
+
+		auto &token = tokens[i];
+		size_t start = token.start > start_pos ? token.start : start_pos;
+		size_t end = next > end_pos ? end_pos : next;
+		if (end <= start) {
+			continue;
+		}
+		if (prev_pos > start) {
+#ifdef DEBUG
+			throw InternalException("ERROR - Rendering at position %llu after rendering at position %llu\n", start,
+			                        prev_pos);
+#endif
+			Linenoise::Log("ERROR - Rendering at position %llu after rendering at position %llu\n", start, prev_pos);
+			continue;
+		}
+		prev_pos = start;
+		std::string text = std::string(buf + start, end - start);
+		if (token.search_match) {
+			ss << underline;
+		}
+		switch (token.type) {
+		case tokenType::TOKEN_KEYWORD:
+			ss << keyword << text << reset;
+			break;
+		case tokenType::TOKEN_NUMERIC_CONSTANT:
+		case tokenType::TOKEN_STRING_CONSTANT:
+			ss << constant << text << reset;
+			break;
+		case tokenType::TOKEN_CONTINUATION:
+			ss << continuation << text << reset;
+			break;
+		case tokenType::TOKEN_CONTINUATION_SELECTED:
+			ss << continuation_selected << text << reset;
+			break;
+		case tokenType::TOKEN_BRACKET:
+			ss << underline << text << reset;
+			break;
+		case tokenType::TOKEN_ERROR:
+			ss << error << text << reset;
+			break;
+		case tokenType::TOKEN_COMMENT:
+			ss << comment << text << reset;
+			break;
+		default:
+			ss << text;
+			if (token.search_match) {
+				ss << reset;
+			}
+		}
+	}
+	return ss.str();
+}
+
+} // namespace duckdb