Move tree-sitter code index updating to an async process to unblock the ui while processing large files

This commit is contained in:
Peter Slattery 2025-07-14 11:22:39 -07:00
parent ae7440aa0b
commit 3f9b803c62
2 changed files with 109 additions and 113 deletions

View File

@ -103,6 +103,34 @@ tree_sitter_language_for_buffer(Application_Links* app, Buffer_ID buffer_id)
return tree_sitter_language_for_buffer(app, buffer_id, scratch);
}
////////////////////////////////////////////////////////////////////
// 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);
}
/////////////////////////////////////////////
// Tree Sitter Hook Internals
/////////////////////////////////////////////
@ -264,10 +292,10 @@ tree_sitter_parse_async__inner(Async_Context* actx, Buffer_ID buffer_id)
tree_data->tree = new_tree;
system_mutex_acquire(tree_data->tree_mutex);
Tree_Sitter_Code_Index_Parse_State parse_state = tree_sitter_code_index_update_single_buffer(app, buffer_id);
code_index_lock();
code_index_set_file(parse_state.buffer_id, parse_state.index_arena, parse_state.index);
code_index_unlock();
Managed_Scope scope = buffer_get_managed_scope(app, buffer_id);
Async_Task* lex_task = scope_attachment(app, scope, buffer_tree_sitter_update_code_index_task_id, Async_Task);
if (lex_task != 0) async_task_cancel_nowait(app, &global_async_system, *lex_task);
*lex_task = async_task_no_dep(&global_async_system, tree_sitter_code_index_update_async, make_data_struct(&buffer_id));
// Force a frame refresh by requesting another frame
animate_in_n_milliseconds(app, 0);
@ -289,6 +317,10 @@ tree_sitter_parse_async(Async_Context* actx, String_Const_u8 data)
tree_sitter_parse_async__inner(actx, buffer_id);
}
////////////////////////////////////////////////////////////////////
// Code Index Updating
////////////////////////////////////////////////////////////////////
function Range_i64
tree_sitter_node_to_range(TSNode node)
{
@ -312,8 +344,11 @@ code_index_new_note(Code_Index_File* index, Arena* arena, Code_Index_Note_Kind k
}
function Tree_Sitter_Code_Index_Parse_State
tree_sitter_code_index_update_single_buffer(Application_Links* app, Buffer_ID buffer_id)
tree_sitter_code_index_update_single_buffer(Async_Context* actx, Buffer_ID buffer_id)
{
Application_Links* app = actx->app;
acquire_global_frame_mutex(app);
Scratch_Block scratch(app);
Tree_Sitter_Code_Index_Parse_State parse_state = {};
@ -321,7 +356,6 @@ tree_sitter_code_index_update_single_buffer(Application_Links* app, Buffer_ID bu
parse_state.buffer_id = buffer_id;
parse_state.language = tree_sitter_language_for_buffer(app, parse_state.buffer_id);
if (!parse_state.language) return parse_state;
parse_state.index_arena = make_arena_system(KB(16));
parse_state.index = push_array_zero(&parse_state.index_arena, Code_Index_File, 1);
parse_state.buffer_contents = push_whole_buffer(app, scratch, parse_state.buffer_id);
@ -332,10 +366,12 @@ tree_sitter_code_index_update_single_buffer(Application_Links* app, Buffer_ID bu
parse_state.nest_stack_last = 0;
parse_state.last_note = 0;
parse_state.last_note_match_id = max_u32;
release_global_frame_mutex(app);
parse_state.ok = true;
TSQueryMatch query_match;
u32 capture_index;
while (tree_sitter_query_continue(&parse_state.query, &query_match, &capture_index))
while (parse_state.ok && tree_sitter_query_continue(&parse_state.query, &query_match, &capture_index))
{
TSQueryCapture type_capture = query_match.captures[capture_index];
@ -408,8 +444,18 @@ tree_sitter_code_index_update_single_buffer(Application_Links* app, Buffer_ID bu
parse_state.last_note->pos = Ii64_size(type_range.start, type_range.end - type_range.start);
}
}
if (async_check_canceled(actx))
{
parse_state.ok = false;
break;
}
}
tree_sitter_query_end(&parse_state.query);
if (parse_state.ok)
{
while (parse_state.nest_stack_last != 0)
{
Code_Index_Nest* nest = parse_state.nest_stack_last->nest;
@ -428,96 +474,44 @@ tree_sitter_code_index_update_single_buffer(Application_Links* app, Buffer_ID bu
parse_state.index->nest_array = code_index_nest_ptr_array_from_list(&parse_state.index_arena, &parse_state.index->nest_list);
parse_state.index->note_array = code_index_note_ptr_array_from_list(&parse_state.index_arena, &parse_state.index->note_list);
tree_sitter_query_end(&parse_state.query);
code_index_lock();
code_index_set_file(parse_state.buffer_id, parse_state.index_arena, parse_state.index);
code_index_unlock();
buffer_clear_layout_cache(app, parse_state.buffer_id);
}
else
{
linalloc_clear(&parse_state.index_arena);
}
parse_state.ok = true;
return parse_state;
}
#define BUFFER_CODE_INDEX_UPDATES false
function void
tree_sitter_code_index_update_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_code_index_update_single_buffer(actx, buffer_id);
}
function void
tree_sitter_code_index_update_tick(Application_Links* app)
{
#define BUFFERED_CODE_INDICES_CAP 8
Tree_Sitter_Code_Index_Parse_State buffered_code_indices[BUFFERED_CODE_INDICES_CAP];
u32 buffered_code_indices_count = 0;
Tree_Sitter_Code_Index_Parse_State parse_state = {};
for (Buffer_Modified_Node* modified_node = global_buffer_modified_set.first;
modified_node != 0;
modified_node = modified_node->next
){
parse_state = tree_sitter_code_index_update_single_buffer(app, modified_node->buffer);
if (parse_state.ok)
{
#if BUFFER_CODE_INDEX_UPDATES
buffered_code_indices[buffered_code_indices_count] = parse_state;
buffered_code_indices_count += 1;
#else
code_index_lock();
code_index_set_file(parse_state.buffer_id, parse_state.index_arena, parse_state.index);
code_index_unlock();
#endif
Buffer_ID buffer_id = modified_node->buffer;
Managed_Scope scope = buffer_get_managed_scope(app, buffer_id);
Async_Task* lex_task = scope_attachment(app, scope, buffer_tree_sitter_update_code_index_task_id, Async_Task);
if (lex_task != 0) async_task_cancel_nowait(app, &global_async_system, *lex_task);
*lex_task = async_task_no_dep(&global_async_system, tree_sitter_code_index_update_async, make_data_struct(&buffer_id));
}
if (buffered_code_indices_count >= BUFFERED_CODE_INDICES_CAP)
{
code_index_lock();
for (int i = 0; i < buffered_code_indices_count; i++)
{
parse_state = buffered_code_indices[i];
code_index_set_file(parse_state.buffer_id, parse_state.index_arena, parse_state.index);
}
buffered_code_indices_count = 0;
code_index_unlock();
}
buffer_clear_layout_cache(app, parse_state.buffer_id);
}
if (buffered_code_indices_count > 0)
{
code_index_lock();
for (int i = 0; i < buffered_code_indices_count; i++)
{
parse_state = buffered_code_indices[i];
code_index_set_file(parse_state.buffer_id, parse_state.index_arena, parse_state.index);
}
code_index_unlock();
}
buffer_modified_set_clear();
}
////////////////////////////////////////////////////////////////////
// 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
////////////////////////////////////////////////////////////////////

View File

@ -92,6 +92,7 @@ global Tree_Sitter_Languages tree_sitter_languages;
CUSTOM_ID(attachment, buffer_tree_sitter_data_id);
CUSTOM_ID(attachment, buffer_tree_sitter_parse_task_id);
CUSTOM_ID(attachment, buffer_tree_sitter_update_code_index_task_id);
b8 use_tree_sitter_code_indexing = true;
b8 use_tree_sitter_token_coloring = true;
@ -105,5 +106,6 @@ function Tree_Sitter_Language_Definition* tree_sitter_language_for_buffer(Applic
function Tree_Sitter_Code_Index_Parse_State tree_sitter_code_index_update_single_buffer(Application_Links *app, Buffer_ID buffer_id);
function void tree_sitter_code_index_update_tick(Application_Links *app);
function void tree_sitter_code_index_update_async(Async_Context* actx, String_Const_u8 data);
#endif //FCODER_TREE_SITTER_H