Codeprocastinator/optimized-tinyllama-covalent
0119
1#include "console.h"2#include <vector>3#include <iostream>4 5#if defined(_WIN32)6#define WIN32_LEAN_AND_MEAN7#ifndef NOMINMAX8#define NOMINMAX9#endif10#include <windows.h>11#include <fcntl.h>12#include <io.h>13#ifndef ENABLE_VIRTUAL_TERMINAL_PROCESSING14#define ENABLE_VIRTUAL_TERMINAL_PROCESSING 0x000415#endif16#else17#include <climits>18#include <sys/ioctl.h>19#include <unistd.h>20#include <wchar.h>21#include <stdio.h>22#include <stdlib.h>23#include <signal.h>24#include <termios.h>25#endif26 27#define ANSI_COLOR_RED "\x1b[31m"28#define ANSI_COLOR_GREEN "\x1b[32m"29#define ANSI_COLOR_YELLOW "\x1b[33m"30#define ANSI_COLOR_BLUE "\x1b[34m"31#define ANSI_COLOR_MAGENTA "\x1b[35m"32#define ANSI_COLOR_CYAN "\x1b[36m"33#define ANSI_COLOR_RESET "\x1b[0m"34#define ANSI_BOLD "\x1b[1m"35 36namespace console {37 38 //39 // Console state40 //41 42 static bool advanced_display = false;43 static bool simple_io = true;44 static display_t current_display = reset;45 46 static FILE* out = stdout;47 48#if defined (_WIN32)49 static void* hConsole;50#else51 static FILE* tty = nullptr;52 static termios initial_state;53#endif54 55 //56 // Init and cleanup57 //58 59 void init(bool use_simple_io, bool use_advanced_display) {60 advanced_display = use_advanced_display;61 simple_io = use_simple_io;62#if defined(_WIN32)63 // Windows-specific console initialization64 DWORD dwMode = 0;65 hConsole = GetStdHandle(STD_OUTPUT_HANDLE);66 if (hConsole == INVALID_HANDLE_VALUE || !GetConsoleMode(hConsole, &dwMode)) {67 hConsole = GetStdHandle(STD_ERROR_HANDLE);68 if (hConsole != INVALID_HANDLE_VALUE && (!GetConsoleMode(hConsole, &dwMode))) {69 hConsole = nullptr;70 simple_io = true;71 }72 }73 if (hConsole) {74 // Check conditions combined to reduce nesting75 if (advanced_display && !(dwMode & ENABLE_VIRTUAL_TERMINAL_PROCESSING) &&76 !SetConsoleMode(hConsole, dwMode | ENABLE_VIRTUAL_TERMINAL_PROCESSING)) {77 advanced_display = false;78 }79 // Set console output codepage to UTF880 SetConsoleOutputCP(CP_UTF8);81 }82 HANDLE hConIn = GetStdHandle(STD_INPUT_HANDLE);83 if (hConIn != INVALID_HANDLE_VALUE && GetConsoleMode(hConIn, &dwMode)) {84 // Set console input codepage to UTF1685 _setmode(_fileno(stdin), _O_WTEXT);86 87 // Set ICANON (ENABLE_LINE_INPUT) and ECHO (ENABLE_ECHO_INPUT)88 if (simple_io) {89 dwMode |= ENABLE_LINE_INPUT | ENABLE_ECHO_INPUT;90 } else {91 dwMode &= ~(ENABLE_LINE_INPUT | ENABLE_ECHO_INPUT);92 }93 if (!SetConsoleMode(hConIn, dwMode)) {94 simple_io = true;95 }96 }97 if (simple_io) {98 _setmode(_fileno(stdin), _O_U8TEXT);99 }100#else101 // POSIX-specific console initialization102 if (!simple_io) {103 struct termios new_termios;104 tcgetattr(STDIN_FILENO, &initial_state);105 new_termios = initial_state;106 new_termios.c_lflag &= ~(ICANON | ECHO);107 new_termios.c_cc[VMIN] = 1;108 new_termios.c_cc[VTIME] = 0;109 tcsetattr(STDIN_FILENO, TCSANOW, &new_termios);110 111 tty = fopen("/dev/tty", "w+");112 if (tty != nullptr) {113 out = tty;114 }115 }116 117 setlocale(LC_ALL, "");118#endif119 }120 121 void cleanup() {122 // Reset console display123 set_display(reset);124 125#if !defined(_WIN32)126 // Restore settings on POSIX systems127 if (!simple_io) {128 if (tty != nullptr) {129 out = stdout;130 fclose(tty);131 tty = nullptr;132 }133 tcsetattr(STDIN_FILENO, TCSANOW, &initial_state);134 }135#endif136 }137 138 //139 // Display and IO140 //141 142 // Keep track of current display and only emit ANSI code if it changes143 void set_display(display_t display) {144 if (advanced_display && current_display != display) {145 fflush(stdout);146 switch(display) {147 case reset:148 fprintf(out, ANSI_COLOR_RESET);149 break;150 case prompt:151 fprintf(out, ANSI_COLOR_YELLOW);152 break;153 case user_input:154 fprintf(out, ANSI_BOLD ANSI_COLOR_GREEN);155 break;156 case error:157 fprintf(out, ANSI_BOLD ANSI_COLOR_RED);158 }159 current_display = display;160 fflush(out);161 }162 }163 164 static char32_t getchar32() {165#if defined(_WIN32)166 HANDLE hConsole = GetStdHandle(STD_INPUT_HANDLE);167 wchar_t high_surrogate = 0;168 169 while (true) {170 INPUT_RECORD record;171 DWORD count;172 if (!ReadConsoleInputW(hConsole, &record, 1, &count) || count == 0) {173 return WEOF;174 }175 176 if (record.EventType == KEY_EVENT && record.Event.KeyEvent.bKeyDown) {177 wchar_t wc = record.Event.KeyEvent.uChar.UnicodeChar;178 if (wc == 0) {179 continue;180 }181 182 if ((wc >= 0xD800) && (wc <= 0xDBFF)) { // Check if wc is a high surrogate183 high_surrogate = wc;184 continue;185 }186 if ((wc >= 0xDC00) && (wc <= 0xDFFF)) { // Check if wc is a low surrogate187 if (high_surrogate != 0) { // Check if we have a high surrogate188 return ((high_surrogate - 0xD800) << 10) + (wc - 0xDC00) + 0x10000;189 }190 }191 192 high_surrogate = 0; // Reset the high surrogate193 return static_cast<char32_t>(wc);194 }195 }196#else197 wchar_t wc = getwchar();198 if (static_cast<wint_t>(wc) == WEOF) {199 return WEOF;200 }201 202#if WCHAR_MAX == 0xFFFF203 if ((wc >= 0xD800) && (wc <= 0xDBFF)) { // Check if wc is a high surrogate204 wchar_t low_surrogate = getwchar();205 if ((low_surrogate >= 0xDC00) && (low_surrogate <= 0xDFFF)) { // Check if the next wchar is a low surrogate206 return (static_cast<char32_t>(wc & 0x03FF) << 10) + (low_surrogate & 0x03FF) + 0x10000;207 }208 }209 if ((wc >= 0xD800) && (wc <= 0xDFFF)) { // Invalid surrogate pair210 return 0xFFFD; // Return the replacement character U+FFFD211 }212#endif213 214 return static_cast<char32_t>(wc);215#endif216 }217 218 static void pop_cursor() {219#if defined(_WIN32)220 if (hConsole != NULL) {221 CONSOLE_SCREEN_BUFFER_INFO bufferInfo;222 GetConsoleScreenBufferInfo(hConsole, &bufferInfo);223 224 COORD newCursorPosition = bufferInfo.dwCursorPosition;225 if (newCursorPosition.X == 0) {226 newCursorPosition.X = bufferInfo.dwSize.X - 1;227 newCursorPosition.Y -= 1;228 } else {229 newCursorPosition.X -= 1;230 }231 232 SetConsoleCursorPosition(hConsole, newCursorPosition);233 return;234 }235#endif236 putc('\b', out);237 }238 239 static int estimateWidth(char32_t codepoint) {240#if defined(_WIN32)241 (void)codepoint;242 return 1;243#else244 return wcwidth(codepoint);245#endif246 }247 248 static int put_codepoint(const char* utf8_codepoint, size_t length, int expectedWidth) {249#if defined(_WIN32)250 CONSOLE_SCREEN_BUFFER_INFO bufferInfo;251 if (!GetConsoleScreenBufferInfo(hConsole, &bufferInfo)) {252 // go with the default253 return expectedWidth;254 }255 COORD initialPosition = bufferInfo.dwCursorPosition;256 DWORD nNumberOfChars = length;257 WriteConsole(hConsole, utf8_codepoint, nNumberOfChars, &nNumberOfChars, NULL);258 259 CONSOLE_SCREEN_BUFFER_INFO newBufferInfo;260 GetConsoleScreenBufferInfo(hConsole, &newBufferInfo);261 262 // Figure out our real position if we're in the last column263 if (utf8_codepoint[0] != 0x09 && initialPosition.X == newBufferInfo.dwSize.X - 1) {264 DWORD nNumberOfChars;265 WriteConsole(hConsole, &" \b", 2, &nNumberOfChars, NULL);266 GetConsoleScreenBufferInfo(hConsole, &newBufferInfo);267 }268 269 int width = newBufferInfo.dwCursorPosition.X - initialPosition.X;270 if (width < 0) {271 width += newBufferInfo.dwSize.X;272 }273 return width;274#else275 // We can trust expectedWidth if we've got one276 if (expectedWidth >= 0 || tty == nullptr) {277 fwrite(utf8_codepoint, length, 1, out);278 return expectedWidth;279 }280 281 fputs("\033[6n", tty); // Query cursor position282 int x1;283 int y1;284 int x2;285 int y2;286 int results = 0;287 results = fscanf(tty, "\033[%d;%dR", &y1, &x1);288 289 fwrite(utf8_codepoint, length, 1, tty);290 291 fputs("\033[6n", tty); // Query cursor position292 results += fscanf(tty, "\033[%d;%dR", &y2, &x2);293 294 if (results != 4) {295 return expectedWidth;296 }297 298 int width = x2 - x1;299 if (width < 0) {300 // Calculate the width considering text wrapping301 struct winsize w;302 ioctl(STDOUT_FILENO, TIOCGWINSZ, &w);303 width += w.ws_col;304 }305 return width;306#endif307 }308 309 static void replace_last(char ch) {310#if defined(_WIN32)311 pop_cursor();312 put_codepoint(&ch, 1, 1);313#else314 fprintf(out, "\b%c", ch);315#endif316 }317 318 static void append_utf8(char32_t ch, std::string & out) {319 if (ch <= 0x7F) {320 out.push_back(static_cast<unsigned char>(ch));321 } else if (ch <= 0x7FF) {322 out.push_back(static_cast<unsigned char>(0xC0 | ((ch >> 6) & 0x1F)));323 out.push_back(static_cast<unsigned char>(0x80 | (ch & 0x3F)));324 } else if (ch <= 0xFFFF) {325 out.push_back(static_cast<unsigned char>(0xE0 | ((ch >> 12) & 0x0F)));326 out.push_back(static_cast<unsigned char>(0x80 | ((ch >> 6) & 0x3F)));327 out.push_back(static_cast<unsigned char>(0x80 | (ch & 0x3F)));328 } else if (ch <= 0x10FFFF) {329 out.push_back(static_cast<unsigned char>(0xF0 | ((ch >> 18) & 0x07)));330 out.push_back(static_cast<unsigned char>(0x80 | ((ch >> 12) & 0x3F)));331 out.push_back(static_cast<unsigned char>(0x80 | ((ch >> 6) & 0x3F)));332 out.push_back(static_cast<unsigned char>(0x80 | (ch & 0x3F)));333 } else {334 // Invalid Unicode code point335 }336 }337 338 // Helper function to remove the last UTF-8 character from a string339 static void pop_back_utf8_char(std::string & line) {340 if (line.empty()) {341 return;342 }343 344 size_t pos = line.length() - 1;345 346 // Find the start of the last UTF-8 character (checking up to 4 bytes back)347 for (size_t i = 0; i < 3 && pos > 0; ++i, --pos) {348 if ((line[pos] & 0xC0) != 0x80) {349 break; // Found the start of the character350 }351 }352 line.erase(pos);353 }354 355 static bool readline_advanced(std::string & line, bool multiline_input) {356 if (out != stdout) {357 fflush(stdout);358 }359 360 line.clear();361 std::vector<int> widths;362 bool is_special_char = false;363 bool end_of_stream = false;364 365 char32_t input_char;366 while (true) {367 fflush(out); // Ensure all output is displayed before waiting for input368 input_char = getchar32();369 370 if (input_char == '\r' || input_char == '\n') {371 break;372 }373 374 if (input_char == (char32_t) WEOF || input_char == 0x04 /* Ctrl+D*/) {375 end_of_stream = true;376 break;377 }378 379 if (is_special_char) {380 set_display(user_input);381 replace_last(line.back());382 is_special_char = false;383 }384 385 if (input_char == '\033') { // Escape sequence386 char32_t code = getchar32();387 if (code == '[' || code == 0x1B) {388 // Discard the rest of the escape sequence389 while ((code = getchar32()) != (char32_t) WEOF) {390 if ((code >= 'A' && code <= 'Z') || (code >= 'a' && code <= 'z') || code == '~') {391 break;392 }393 }394 }395 } else if (input_char == 0x08 || input_char == 0x7F) { // Backspace396 if (!widths.empty()) {397 int count;398 do {399 count = widths.back();400 widths.pop_back();401 // Move cursor back, print space, and move cursor back again402 for (int i = 0; i < count; i++) {403 replace_last(' ');404 pop_cursor();405 }406 pop_back_utf8_char(line);407 } while (count == 0 && !widths.empty());408 }409 } else {410 int offset = line.length();411 append_utf8(input_char, line);412 int width = put_codepoint(line.c_str() + offset, line.length() - offset, estimateWidth(input_char));413 if (width < 0) {414 width = 0;415 }416 widths.push_back(width);417 }418 419 if (!line.empty() && (line.back() == '\\' || line.back() == '/')) {420 set_display(prompt);421 replace_last(line.back());422 is_special_char = true;423 }424 }425 426 bool has_more = multiline_input;427 if (is_special_char) {428 replace_last(' ');429 pop_cursor();430 431 char last = line.back();432 line.pop_back();433 if (last == '\\') {434 line += '\n';435 fputc('\n', out);436 has_more = !has_more;437 } else {438 // llama will just eat the single space, it won't act as a space439 if (line.length() == 1 && line.back() == ' ') {440 line.clear();441 pop_cursor();442 }443 has_more = false;444 }445 } else {446 if (end_of_stream) {447 has_more = false;448 } else {449 line += '\n';450 fputc('\n', out);451 }452 }453 454 fflush(out);455 return has_more;456 }457 458 static bool readline_simple(std::string & line, bool multiline_input) {459#if defined(_WIN32)460 std::wstring wline;461 if (!std::getline(std::wcin, wline)) {462 // Input stream is bad or EOF received463 line.clear();464 GenerateConsoleCtrlEvent(CTRL_C_EVENT, 0);465 return false;466 }467 468 int size_needed = WideCharToMultiByte(CP_UTF8, 0, &wline[0], (int)wline.size(), NULL, 0, NULL, NULL);469 line.resize(size_needed);470 WideCharToMultiByte(CP_UTF8, 0, &wline[0], (int)wline.size(), &line[0], size_needed, NULL, NULL);471#else472 if (!std::getline(std::cin, line)) {473 // Input stream is bad or EOF received474 line.clear();475 return false;476 }477#endif478 if (!line.empty()) {479 char last = line.back();480 if (last == '/') { // Always return control on '/' symbol481 line.pop_back();482 return false;483 }484 if (last == '\\') { // '\\' changes the default action485 line.pop_back();486 multiline_input = !multiline_input;487 }488 }489 line += '\n';490 491 // By default, continue input if multiline_input is set492 return multiline_input;493 }494 495 bool readline(std::string & line, bool multiline_input) {496 set_display(user_input);497 498 if (simple_io) {499 return readline_simple(line, multiline_input);500 }501 return readline_advanced(line, multiline_input);502 }503 504}505