4coder/code/custom/4coder_tree_sitter.cpp

813 lines
22 KiB
C++

/////////////////////////////////////////////
// TEMP until I implement more generic language stuff
/////////////////////////////////////////////
/////////////////////////////////////////////
// C
/////////////////////////////////////////////
// C++
TSQuery* tree_sitter_cpp_index_query;
String_Const_u8 TS_CPP_INDEX_QUERY = string_u8_litexpr("(_ \"{\" @Start \"}\" @End ) @ScopeNest\n");
String_Const_u8 TS_CPP_HIGHLIGHT_QUERY = string_u8_litexpr(R"DONE(
(call_expression function: [
(identifier) @defcolor_function
(field_expression field: (field_identifier) @defcolor_function)])
(function_declarator
declarator: [(identifier) (field_identifier)] @defcolor_function)
(preproc_def
name: (identifier) @defcolor_macro)
(preproc_function_def
name: (identifier) @defcolor_macro)
(type_identifier) @defcolor_type
(call_expression
function: (parenthesized_expression
(identifier) @defcolor_type))
[(primitive_type) (type_qualifier) (storage_class_specifier)
(break_statement) (continue_statement) "union" "return" "do"
"while" "for" "if" "class" "struct" "enum" "sizeof"
"else" "switch" "case"] @defcolor_keyword
[(number_literal) (string_literal) (raw_string_literal)] @defcolor_str_constant
[(preproc_directive) "#define" "#if" "#elif" "#else" "#endif"
"#include"] @defcolor_preproc
["{" "}" ";" ":" ","] @defcolor_text_default
(comment) @defcolor_comment
)DONE");
/////////////////////////////////////////////
// Jai
String_Const_u8 TS_JAI_HIGHLIGHT_QUERY = string_u8_litexpr(R"DONE(
; keywords
[
"if"
"else"
"break"
"continue"
"return"
"struct"
"enum"
"for"
"defer"
"cast"
"xx"
"ifx"
"null"
] @defcolor_keyword
; # preceeded
[
(compiler_directive)
(import)
(char_string)
] @defcolor_macro
(import (identifier) @defcolor_type)
; Identifiers
(struct_declaration
name: (identifier) @defcolor_type
)
(struct_literal
type: (identifier) @defcolor_type
)
(enum_declaration
name: (identifier) @defcolor_type
)
(enum_declaration "{" (identifier) @defcolor_type)
(variable_declaration
type: (types) @defcolor_type
)
(procedure_declaration
name: (identifier) @defcolor_function
)
(call_expression
function: (identifier) @defcolor_function
)
(procedure
result: (procedure_returns) @defcolor_type
)
; Constants & Literals
[
(string)
(string_directive)
] @defcolor_str_constant
(escape_sequence) @defcolor_special_character
(integer) @defcolor_int_constant
(float) @defcolor_float_constant
(boolean) @defcolor_bool_constant
(array_literal
type: (identifier) @defcolor_type
)
; Comments
(note) @defcolor_comment
(block_comment) @defcolor_comment
(block_comment_text) @defcolor_comment
)DONE");
// NOTE(PS): source: https://github.com/St0wy/tree-sitter-jai/blob/main/queries/highlights.scm
String_Const_u8 TS_JAI_HIGHLIGHT_QUERY_ = string_u8_litexpr(R"DONE(
[
(compiler_directive)
(import)
] @defcolor_macro
; Keywords
; TODO : complete this list
[
"struct"
"enum"
"defer"
"cast"
"xx"
"return"
] @defcolor_keyword
; Conditionals
[
"if"
"else"
"case"
"break"
] @defcolor_keyword
((if_expression
[
"then"
"ifx"
"else"
] @defcolor_keyword)
(#set! "priority" 105))
; Repeats
[
"for"
"while"
"continue"
] @defcolor_keyword
; Variables
(identifier) @defcolor_text_default
; Namespaces
(import (identifier) @defcolor_text_default)
; Parameters
(parameter (identifier) @defcolor_text_default ":" "="? (identifier)? @defcolor_str_constant)
(default_parameter (identifier) @defcolor_text_default ":=")
(call_expression argument: (identifier) @defcolor_text_default "=")
; Functions
(procedure_declaration (identifier) @defcolor_function)
(procedure_declaration (identifier) @defcolor_function (procedure (block)))
(call_expression function: (identifier) @defcolor_function)
; Types
(type (identifier) @defcolor_type)
((type (identifier) @defcolor_type)
(#any-of? @type.builtin
"bool"
"int" "s8" "s16" "s32" "s64"
"u8" "u16" "u32" "u64"
"string"))
(struct_declaration (identifier) @defcolor_type "::")
(enum_declaration (identifier) @defcolor_type "::")
;(union_declaration (identifier) @defcolor_type "::")
(const_declaration (identifier) @defcolor_type "::" [(array_type) (pointer_type)])
(struct . (identifier) @defcolor_type)
;(field_type . (identifier) @namespace "." (identifier) @defcolor_type)
;(bit_set_type (identifier) @defcolor_type ";")
;(procedure_type (parameters (parameter (identifier) @defcolor_type)))
;(polymorphic_parameters (identifier) @defcolor_type)
((identifier) @defcolor_type
(#lua-match? @defcolor_type "^[A-Z][a-zA-Z0-9]*$")
(#not-has-parent? @defcolor_type parameter procedure_declaration call_expression))
; Fields
(member_expression "." (identifier) @defcolor_text_default)
;(struct_type "{" (identifier) @defcolor_text_default)
(struct_field (identifier) @defcolor_text_default "="?)
(field (identifier) @defcolor_text_default)
; Constants
((identifier) @defcolor_text_default
(#lua-match? @defcolor_str_constnat "^_*[A-Z][A-Z0-9_]*$")
(#not-has-parent? @text_default type parameter))
(member_expression . "." (identifier) @defcolor_text_default)
(enum_declaration "{" (identifier) @defcolor_text_default)
; Literals
(number) @defcolor_int_constant
(float) @defcolor_float_constant
(string) @defcolor_str_constnat
;(character) @defcolor_str_constnat
(escape_sequence) @defcolor_str_constant
(boolean) @defcolor_bool_constant
[
(uninitialized)
(null)
] @defcolor_text_default
((identifier) @defcolor_text_default
(#any-of? @defcolor_text_default "context"))
; Operators
[
":="
"="
"+"
"-"
"*"
"/"
"%"
"%%"
">"
">="
"<"
"<="
"=="
"!="
"~="
"|"
"~"
"&"
"&~"
"<<"
">>"
"||"
"&&"
"!"
".."
"+="
"-="
"*="
"/="
"%="
"&="
"|="
"^="
"<<="
">>="
"||="
"&&="
"&~="
;"..="
;"..<"
;"?"
] @defcolor_operator
; Punctuation
[ "{" "}" ] @defcolor_text_default
[ "(" ")" ] @defcolor_text_default
[ "[" "]" ] @defcolor_text_default
[
"::"
"->"
"."
","
":"
";"
] @defcolor_text_default
; Comments
[
(comment)
(block_comment)
] @defcolor_comment
; Errors
(ERROR) @defcolor_comment_pop
)DONE");
////////////////////////////////////////////////////////////////////
// Language Management
////////////////////////////////////////////////////////////////////
function void
tree_sitter_register_language(String_Const_u8 ext, TSLanguage* language, TSQuery* highlight_query)
{
Tree_Sitter_Language_Definition* lang = 0;
u64 hash = table_hash_u8(ext.str, ext.size);
u64 slot = hash % ArrayCount(tree_sitter_languages.languages);
for (Tree_Sitter_Language_Definition* l = tree_sitter_languages.languages[slot]; l != 0; l = l->next)
{
if (l->extension_hash == hash && string_match(l->extension, ext))
{
lang = l; break;
}
}
if (lang == 0)
{
lang = push_array(&tree_sitter_languages.arena, Tree_Sitter_Language_Definition, 1);
lang->next = tree_sitter_languages.languages[slot];
tree_sitter_languages.languages[slot] = lang;
lang->extension_hash = hash;
lang->extension = push_string_copy(&tree_sitter_languages.arena, ext);
lang->language = language;
lang->highlight_query = highlight_query;
}
}
function Tree_Sitter_Language_Definition*
tree_sitter_language_from_file_extension(String_Const_u8 ext)
{
Tree_Sitter_Language_Definition* result = 0;
u64 ext_hash = table_hash_u8(ext.str, ext.size);
u64 slot = ext_hash % ArrayCount(tree_sitter_languages.languages);
for (Tree_Sitter_Language_Definition* l = tree_sitter_languages.languages[slot]; l != 0; l = l->next)
{
if (l->extension_hash == ext_hash && string_match(l->extension, ext))
{
result = l;
break;
}
}
return result;
}
function Tree_Sitter_Language_Definition*
tree_sitter_language_for_buffer(Application_Links* app, Buffer_ID buffer_id, Arena* arena)
{
Tree_Sitter_Language_Definition* result = 0;
String_Const_u8 file_name = push_buffer_file_name(app, arena, buffer_id);
String_Const_u8 extension = string_file_extension(file_name);
result = tree_sitter_language_from_file_extension(extension);
return result;
}
function Tree_Sitter_Language_Definition*
tree_sitter_language_for_buffer(Application_Links* app, Buffer_ID buffer_id)
{
Scratch_Block scratch(app);
return tree_sitter_language_for_buffer(app, buffer_id, scratch);
}
/////////////////////////////////////////////
// Tree Sitter Hook Internals
/////////////////////////////////////////////
function void
tree_sitter_init(Application_Links* app)
{
Buffer_ID buffer = create_buffer(
app,
string_u8_litexpr("*tree*"),
BufferCreate_NeverAttachToFile | BufferCreate_AlwaysNew
);
buffer_set_setting(app, buffer, BufferSetting_Unimportant, true);
buffer_set_setting(app, buffer, BufferSetting_ReadOnly, true);
tree_sitter_languages.arena = make_arena_system(KB(16));
u32 error_offset;
TSQueryError query_error;
{ // Register CPP
TSLanguage* language = tree_sitter_cpp();
String_Const_u8 highlight_query_str = TS_CPP_HIGHLIGHT_QUERY;
TSQuery* highlight_query = ts_query_new(
language,
(const char *)highlight_query_str.str,
(u32)highlight_query_str.size,
&error_offset, &query_error
);
tree_sitter_register_language(SCu8("c"), language, highlight_query);
tree_sitter_register_language(SCu8("cpp"), language, highlight_query);
tree_sitter_register_language(SCu8("h"), language, highlight_query);
tree_sitter_register_language(SCu8("hpp"), language, highlight_query);
tree_sitter_register_language(SCu8("cc"), language, highlight_query);
}
{ // Register Jai
TSLanguage* language = tree_sitter_jai();
String_Const_u8 highlight_query_str = TS_JAI_HIGHLIGHT_QUERY;
TSQuery* highlight_query = ts_query_new(
language,
(const char *)highlight_query_str.str,
(u32)highlight_query_str.size,
&error_offset, &query_error
);
if (!highlight_query)
{
print_message(app, SCu8("Error creating jai highlight query\n"));
printf("%.*s\n", (int)Min(highlight_query_str.size-error_offset, 100), highlight_query_str.str + error_offset);
}
else
{
tree_sitter_register_language(SCu8("jai"), language, highlight_query);
}
}
}
function void
tree_sitter_begin_buffer(Application_Links* app, Buffer_ID buffer_id)
{
Managed_Scope buffer_scope = buffer_get_managed_scope(app, buffer_id);
Buffer_Tree_Sitter_Data* tree_data = scope_attachment(app, buffer_scope, buffer_tree_sitter_data_id, Buffer_Tree_Sitter_Data);
tree_data->tree_mutex = system_mutex_make();
}
function void
tree_sitter_end_buffer(Application_Links* app, Buffer_ID buffer_id)
{
Managed_Scope buffer_scope = buffer_get_managed_scope(app, buffer_id);
Async_Task *tree_sitter_parse_task = scope_attachment(app, buffer_scope, buffer_tree_sitter_parse_task_id, Async_Task);
if (tree_sitter_parse_task && async_task_is_running_or_pending(&global_async_system, *tree_sitter_parse_task))
{
async_task_cancel(app, &global_async_system, *tree_sitter_parse_task);
}
Buffer_Tree_Sitter_Data* tree_data = scope_attachment(app, buffer_scope, buffer_tree_sitter_data_id, Buffer_Tree_Sitter_Data);
if (!tree_data || !tree_data->tree) return;
system_mutex_acquire(tree_data->tree_mutex);
ts_tree_delete(tree_data->tree);
system_mutex_release(tree_data->tree_mutex);
system_mutex_free(tree_data->tree_mutex);
}
function TSTree*
tree_sitter_buffer_get_tree_copy(Buffer_Tree_Sitter_Data* tree_data)
{
TSTree* result = 0;
// system_mutex_acquire(tree_data->tree_mutex);
if (tree_data->tree) result = ts_tree_copy(tree_data->tree);
// system_mutex_release(tree_data->tree_mutex);
return result;
}
function void
tree_sitter_parse_async__inner(Async_Context* actx, Buffer_ID buffer_id)
{
Application_Links *app = actx->app;
Arena arena = make_arena_system(KB(16));
TSParser *parser = ts_parser_new();
ts_parser_set_timeout_micros(parser, 5000);
acquire_global_frame_mutex(app);
Tree_Sitter_Language_Definition* lang = tree_sitter_language_for_buffer(app, buffer_id);
String_Const_u8 src = push_whole_buffer(app, &arena, buffer_id);
Managed_Scope scope = buffer_get_managed_scope(app, buffer_id);
Buffer_Tree_Sitter_Data* tree_data = scope_attachment(app, scope, buffer_tree_sitter_data_id, Buffer_Tree_Sitter_Data);
TSTree *old_tree = tree_sitter_buffer_get_tree_copy(tree_data);
bool lang_set = ts_parser_set_language(parser, lang->language);
release_global_frame_mutex(app);
if (!lang_set)
{
AssertMessageAlways("Failed to set the language for the parser."
"This probably means a language wasn't set"
"in the BeginBuffer hook.\n");
}
// Iterate until we get a tree or we find that we should cancel the parse
TSTree *new_tree = 0;
b32 canceled = false;
for (;;)
{
new_tree = ts_parser_parse_string(parser, old_tree, (char *)src.str, (u32)src.size);
if (async_check_canceled(actx))
{
canceled = true;
break;
}
if (new_tree) break;
}
if (!canceled && new_tree)
{
TSTree* old_buffer_tree;
acquire_global_frame_mutex(app);
{
// NOTE(jack): Copy the old pointer to delete it outside the mutex.
system_mutex_acquire(tree_data->tree_mutex);
old_buffer_tree = tree_data->tree;
tree_data->tree = new_tree;
system_mutex_acquire(tree_data->tree_mutex);
print_message(app, SCu8("Finished Parse\n"));
// TODO(PS): Just put the code index update call here
// NOTE(jack): This feels kinda hacky, this is here to trigger
// the code index update tick. The buffer is also makred by the
// async lexer so we will update the index too frequently. We
// should probably change the lexer to not mark as modified.
// TODO(jack): Should we instead trigger another async task here to
// update the code index once this is done?
buffer_mark_as_modified(buffer_id);
// Force a frame refresh by requesting another frame
animate_in_n_milliseconds(app, 0);
}
release_global_frame_mutex(app);
ts_tree_delete(old_buffer_tree);
}
ts_parser_delete(parser);
ts_tree_delete(old_tree);
linalloc_clear(&arena);
}
function void
tree_sitter_parse_async(Async_Context* actx, String_Const_u8 data)
{
if (data.size != sizeof(Buffer_ID)) return;
Buffer_ID buffer_id = *(Buffer_ID*)data.str;
tree_sitter_parse_async__inner(actx, buffer_id);
}
function Range_i64
tree_sitter_node_to_range(TSNode node)
{
Range_i64 result;
result.start = ts_node_start_byte(node);
result.end = ts_node_end_byte(node);
return result;
}
function void
tree_sitter_code_index_update_tick(Application_Links* app)
{
Scratch_Block scratch(app);
#if 0
// TODO(PS): this should be done when we register the language
if (!tree_sitter_cpp_index_query)
{
u32 error_offset;
TSQueryError query_error;
tree_sitter_cpp_index_query = ts_query_new(
tree_sitter_cpp(), TS_CPP_INDEX_QUERY, (u32)TS_CPP_INDEX_QUERY.size, &error_offset, &query_error
);
if (!tree_sitter_cpp_index_query)
{
print_message(app, string_u8_litexpr("Failed to create cpp index query\n");
}
}
for (Buffer_Modified_Node* modified_node = global_buffer_modified_set.first;
modified_node != 0;
modified_node = modified_node->next
){
Temp_Memory_Block temp(scratch);
Buffer_ID buffer_id = modified_node->buffer;
Arena arena = make_arena_system(KB(16));
Code_Index_File* index = push_array_zero(&arena, Code_Index_File, 1);
index->buffer = buffer_id;
Tree_Sitter_Code_Index_Nest_List nests = {};
Managed_Scope buffer_scope = buffer_get_managed_scope(app, buffer_id);
Buffer_Tree_Sitter_Data* tree_data = scope_attachment(app, buffer_scope, buffer_tree_sitter_data_id, Buffer_Tree_Sitter_Data);
TSTree* tree = tree_sitter_buffer_get_tree_copy(tree_data);
if (tree)
{
TSQueryCursor* query_cursor = ts_query_cursor_new();
ts_query_cursor_exec(query_cursor, tree_sitter_cpp_index_query, ts_tree_root_node(tree));
TSQueryMatch query_match;
while (ts_query_cursor_next_match(query_cursor, &match))
{
TSQueryCapture type_capture = match.captures[0];
TSNode type_node = type_capture.node;
Range_i64 type_range = tree_sitter_node_to_range(type_node);
// TODO(PS): PICK UP HERE
}
}
ts_tree_delete(tree);
}
#endif
}
////////////////////////////////////////////////////////////////////
// Async Management
////////////////////////////////////////////////////////////////////
function void
async_task_cancel_nowait(Application_Links* app, Async_System* async_system, Async_Task task)
{
system_mutex_acquire(async_system->mutex);
Async_Node *node = async_get_pending_node(async_system, task);
if (node != 0)
{
dll_remove(&node->node);
async_system->task_count -= 1;
async_free_node(async_system, node);
}
else
{
node = async_get_running_node(async_system, task);
if (node != 0)
{
b32 *cancel_signal = &node->thread->cancel_signal;
atomic_write_b32(cancel_signal, true);
// async_task_wait__inner(app, async_system, task);
}
}
system_mutex_release(async_system->mutex);
}
////////////////////////////////////////////////////////////////////
// Token Highlighting
////////////////////////////////////////////////////////////////////
function void
tree_sitter_highlight_node(
Application_Links* app,
TSQuery* query,
TSNode top_node,
TSQueryCursor* query_cursor,
Text_Layout_ID text_layout_id
){
ts_query_cursor_exec(query_cursor, query, top_node);
TSQueryMatch query_match;
u32 capture_index;
while (ts_query_cursor_next_capture(query_cursor, &query_match, &capture_index))
{
TSQueryCapture capture = query_match.captures[capture_index];
TSNode node = capture.node;
u32 length;
const char* tmp = ts_query_capture_name_for_id(query, capture.index, &length);
String_Const_u8 capture_name = SCu8((char*)tmp, length);
Range_i64 highlight_range = tree_sitter_node_to_range(node);
Managed_ID color_id = managed_id_get(app, SCu8("colors"), capture_name);
if (color_id != 0)
{
paint_text_color_fcolor(app, text_layout_id, highlight_range, fcolor_id(color_id));
}
}
}
function void
draw_tree_sitter_node_colors(Application_Links* app, Text_Layout_ID text_layout_id, Buffer_ID buffer_id)
{
Tree_Sitter_Language_Definition* lang = tree_sitter_language_for_buffer(app, buffer_id);
TSQuery* query = lang->highlight_query;
Range_i64 visible_range = text_layout_get_visible_range(app, text_layout_id);
Managed_Scope buffer_scope = buffer_get_managed_scope(app, buffer_id);
Buffer_Tree_Sitter_Data* tree_data = scope_attachment(app, buffer_scope, buffer_tree_sitter_data_id, Buffer_Tree_Sitter_Data);
TSTree* tree = tree_sitter_buffer_get_tree_copy(tree_data);
if (tree)
{
TSNode root = ts_tree_root_node(tree);
// Get the smallest node that fully contains the visible range
TSNode visible_container = ts_node_descendant_for_byte_range(root, (u32)visible_range.start, (u32)visible_range.end);
TSQueryCursor* query_cursor = ts_query_cursor_new();
if (ts_node_eq(root, visible_container))
{
TSTreeCursor tree_cursor = ts_tree_cursor_new(visible_container);
if (ts_tree_cursor_goto_first_child_for_byte(&tree_cursor, (u32)visible_range.start) != -1)
{
do {
TSNode node = ts_tree_cursor_current_node(&tree_cursor);
Range_i64 child_range = tree_sitter_node_to_range(node);
if (child_range.start > visible_range.end) break;
tree_sitter_highlight_node(app, query, node, query_cursor, text_layout_id);
} while(ts_tree_cursor_goto_next_sibling(&tree_cursor));
}
else
{
// Pathological case - just highligh the whole document. This is probably bad
tree_sitter_highlight_node(app, query, root, query_cursor, text_layout_id);
}
}
else
{
tree_sitter_highlight_node(app, query, visible_container, query_cursor, text_layout_id);
}
ts_query_cursor_delete(query_cursor);
}
ts_tree_delete(tree);
}
////////////////////////////////////////////////////////////////////
// DEBUG
////////////////////////////////////////////////////////////////////
char* prefix_buffer = " ";
function void
write_tree_sitter_tree_to_buffer__inner(Application_Links *app, Arena *arena, Buffer_ID buffer_id,
TSNode cur_node, i32 level = 0, const char *field="")
{
TSPoint start = ts_node_start_point(cur_node);
TSPoint end = ts_node_end_point(cur_node);
// + 1 on ts positions becuase the first line/column are zero in treesitter,
// but 4coder displays as 1 indexed in the filebar.
String_Const_u8 string = push_stringf(arena, "%.*s%s: %s [%d, %d] - [%d, %d]\n",
level*2, prefix_buffer, field, ts_node_type(cur_node),
start.row + 1, start.column + 1,
end.row + 1, end.column + 1);
buffer_replace_range(app, buffer_id, Ii64(buffer_get_size(app, buffer_id)), string);
u32 child_count = ts_node_child_count(cur_node);
for (u32 i = 0; i < child_count; ++i)
{
TSNode child = ts_node_child(cur_node, i);
if (ts_node_is_named(child))
{
field = ts_node_field_name_for_child(cur_node, i);
if (!field) field = "";
write_tree_sitter_tree_to_buffer__inner(app, arena, buffer_id, child, level + 1, field);
}
}
}
CUSTOM_COMMAND_SIG(tree_sitter_write_tree)
CUSTOM_DOC("Write the current buffer's tree sitter tree to *tree*")
{
Scratch_Block scratch(app);
Buffer_ID out_buffer = get_buffer_by_name(app, string_u8_litexpr("*tree*"), Access_Always);
View_ID view = get_active_view(app, Access_Always);
Buffer_ID buffer = view_get_buffer(app, view, Access_Visible);
Managed_Scope scope = buffer_get_managed_scope(app, buffer);
Buffer_Tree_Sitter_Data *tree_data = scope_attachment(app, scope, buffer_tree_sitter_data_id, Buffer_Tree_Sitter_Data);
if (tree_data->tree)
{
TSNode root = ts_tree_root_node(tree_data->tree);
write_tree_sitter_tree_to_buffer__inner(app, scratch, out_buffer, root);
}
}