tree_sitter_parse_full_file_async no longer updates the code index

- I added this previously, and the need for it has disappeared

code/custom/4coder_custom_hooks.cpp

@@ -35,8 +35,9 @@ BUFFER_HOOK_SIG(custom_begin_buffer){
 
   if (treat_as_code)
   {
-    Async_Task* parse_task = scope_attachment(app, scope, buffer_tree_sitter_parse_task_id, Async_Task);
-    *parse_task = async_task_no_dep(&global_async_system, tree_sitter_parse_async, make_data_struct(&buffer_id));
+    // TODO(PS): @AsyncrifyThis???
+    tree_sitter_parse_incremental(app, buffer_id);
+    tree_sitter_code_index_update_full_file(app, buffer_id);
   }
 
   String_Const_u8 buffer_name = push_buffer_base_name(app, scratch, buffer_id);
@@ -79,6 +80,8 @@ BUFFER_HOOK_SIG(custom_end_buffer){
   Marker_List *list = get_marker_list_for_buffer(buffer_id);
   if (list != 0) delete_marker_list(list);
 
+  // TODO(PS): kill the async task for this buffer
+
   tree_sitter_end_buffer(app, buffer_id);
   default_end_buffer(app, buffer_id);
   return(0);
@@ -89,43 +92,23 @@ BUFFER_HOOK_SIG(custom_end_buffer){
 ///////////////////////////////////////////////////////////////////////////
 
 BUFFER_EDIT_RANGE_SIG(custom_buffer_edit_range){
-  // buffer_id, new_range, original_size
   ProfileScope(app, "default edit range");
 
+  Scratch_Block scratch(app);
+  Managed_Scope scope = buffer_get_managed_scope(app, buffer_id);
+
   Range_i64 old_range = Ii64(old_cursor_range.min.pos, old_cursor_range.max.pos);
 
-  buffer_shift_fade_ranges(buffer_id, old_range.max, (new_range.max - old_range.max));
-
-  {
-    code_index_lock();
-    Code_Index_File *file = code_index_get_file(buffer_id);
-    if (file != 0) code_index_shift(file, old_range, range_size(new_range));
-    code_index_unlock();
-  }
-
-  i64 insert_size = range_size(new_range);
-  i64 text_shift = replace_range_shift(old_range, insert_size);
-
-  Scratch_Block scratch(app);
-
-  Managed_Scope scope = buffer_get_managed_scope(app, buffer_id);
-
+  b8 trigger_code_index_update = false;
   { // Tree Sitter
     Buffer_Tree_Sitter_Data* tree_data = scope_attachment(app, scope, buffer_tree_sitter_data_id, Buffer_Tree_Sitter_Data);
+
+    // TODO(PS): if there's not tree_data, we actually want to block
+    // if there's an async parse or index task so we don't end up with
+    // a race condition when we try and incrementally update the tree
+    // and index
     if (tree_data->tree)
     {
-      Async_Task* tree_sitter_parse_task = scope_attachment(app, scope, buffer_tree_sitter_parse_task_id, Async_Task);
-      if (async_task_is_running_or_pending(&global_async_system, *tree_sitter_parse_task))
-      {
-        // NOTE(jack): We dont want to wait for ts_tree delete, as it is slow for large files
-        // (noticable spikes when it fires in 4coder_base_types.cpp ~200Kb)
-        // And if we restart the next task before it attempts the delete we may be able to
-        // avoid the delete entirely as the reference count may have increased from the next
-        // parse attempt
-        async_task_cancel_nowait(app, &global_async_system, *tree_sitter_parse_task);
-        *tree_sitter_parse_task = 0;
-      }
-
       i64 new_end_line = get_line_number_from_pos(app, buffer_id, new_range.end);
       i64 new_end_pos = new_range.end - get_line_start_pos(app, buffer_id, new_end_line);
 
@@ -148,11 +131,27 @@ BUFFER_EDIT_RANGE_SIG(custom_buffer_edit_range){
       };
 
       ts_tree_edit(tree_data->tree, &edit);
-      *tree_sitter_parse_task = async_task_no_dep(&global_async_system, tree_sitter_parse_async, make_data_struct(&buffer_id));
+      tree_sitter_parse_incremental(app, buffer_id);
+
+      tree_data->last_update_old_range = old_range;
+      tree_data->last_update_new_range = new_range;
+      trigger_code_index_update = true;
     }
   }
 
+  buffer_shift_fade_ranges(buffer_id, old_range.max, (new_range.max - old_range.max));
 
+  {
+    code_index_lock();
+    Code_Index_File *file = code_index_get_file(buffer_id);
+    if (file != 0) code_index_shift(file, old_range, range_size(new_range));
+    code_index_unlock();
+  }
+
+  i64 insert_size = range_size(new_range);
+  i64 text_shift = replace_range_shift(old_range, insert_size);
+
+  // TODO(PS): I think we want to replace this with a dumber tokenizer
   Async_Task *lex_task_ptr = scope_attachment(app, scope, buffer_lex_task, Async_Task);
 
   Base_Allocator *allocator = managed_scope_allocator(app, scope);
@@ -249,6 +248,8 @@ BUFFER_EDIT_RANGE_SIG(custom_buffer_edit_range){
 
   loco_on_buffer_edit(app, buffer_id, old_range, new_range);
 
+  if (trigger_code_index_update) buffer_mark_as_modified(buffer_id);
+
   // no meaning for return
   return(0);
 }

code/custom/4coder_helper.cpp

@@ -61,21 +61,6 @@ get_token_array_from_buffer(Application_Links *app, Buffer_ID buffer){
     return(result);
 }
 
-function Token_Array
-get_token_array_from_buffer_no_wait(Application_Links *app, Buffer_ID buffer)
-{
-    Token_Array result = {};
-    Managed_Scope scope = buffer_get_managed_scope(app, buffer);
-    Async_Task *lex_task_ptr = scope_attachment(app, scope, buffer_lex_task, Async_Task);
-    if (lex_task_ptr == 0){
-      Token_Array *ptr = scope_attachment(app, scope, attachment_tokens, Token_Array);
-      if (ptr != 0){
-          result = *ptr;
-      }
-    }
-    return(result);
-}
-
 ////////////////////////////////
 
 function Buffer_Seek

code/custom/4coder_tree_sitter.cpp

@@ -164,12 +164,6 @@ 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);
@@ -182,9 +176,9 @@ function TSTree*
 tree_sitter_buffer_get_tree_copy(Buffer_Tree_Sitter_Data* tree_data)
 {
   TSTree* result = 0;
-  // system_mutex_acquire(tree_data->tree_mutex);
+  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);
+  system_mutex_release(tree_data->tree_mutex);
   return result;
 }
 
@@ -239,27 +233,26 @@ tree_sitter_query_end(Tree_Sitter_Query_Cursor* cursor)
 // Tree Parsing
 ////////////////////////////////////////////////////////////////////
 
-function void
-tree_sitter_parse_async__inner(Async_Context* actx, Buffer_ID buffer_id)
+#define DO_CODE_INDEX_UPDATE 0
+
+function Tree_Sitter_Parse_State
+tree_sitter_parse_state_create(Application_Links* app, Buffer_ID buffer_id)
 {
-  Application_Links *app = actx->app;
+  Tree_Sitter_Parse_State parse_state = {};
+  parse_state.arena = make_arena_system(KB(16));   // TODO(PS): investigate if we need this or if we should just use the scratch block instead
+  parse_state.parser = ts_parser_new();
+  ts_parser_set_timeout_micros(parse_state.parser, 5000);
 
-  Arena arena = make_arena_system(KB(16));
-
-  TSParser *parser = ts_parser_new();
-  ts_parser_set_timeout_micros(parser, 5000);
-
-  acquire_global_frame_mutex(app);
   Scratch_Block scratch(app);
-    String_Const_u8 buffer_name = push_buffer_unique_name(app, scratch, buffer_id);
-    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);
+  parse_state.buffer_name = push_buffer_unique_name(app, scratch, buffer_id);
+  parse_state.lang = tree_sitter_language_for_buffer(app, buffer_id);  // TODO(PS): this might not need to be stored on parse_state
+  parse_state.buffer_contents = push_whole_buffer(app, &parse_state.arena, buffer_id);
 
+  Managed_Scope scope = buffer_get_managed_scope(app, buffer_id);
+  parse_state.tree_data = scope_attachment(app, scope, buffer_tree_sitter_data_id, Buffer_Tree_Sitter_Data);
+  parse_state.old_tree = tree_sitter_buffer_get_tree_copy(parse_state.tree_data);
+
+  bool lang_set = ts_parser_set_language(parse_state.parser, parse_state.lang->language);
   if (!lang_set)
   {
     AssertMessageAlways("Failed to set the language for the parser."
@@ -267,12 +260,45 @@ tree_sitter_parse_async__inner(Async_Context* actx, Buffer_ID buffer_id)
                         "in the BeginBuffer hook.\n");
   }
 
+  return parse_state;
+}
+
+function void
+tree_sitter_parse_state_destroy(Tree_Sitter_Parse_State* parse_state)
+{
+  ts_parser_delete(parse_state->parser);
+  ts_tree_delete(parse_state->old_tree);
+  linalloc_clear(&parse_state->arena);
+}
+
+function TSTree*
+tree_sitter_parse_complete(Tree_Sitter_Parse_State* parse_state, TSTree* new_tree)
+{
+  // NOTE(jack): Copy the old pointer to delete it outside the mutex.
+  system_mutex_acquire(parse_state->tree_data->tree_mutex);
+  TSTree* old_buffer_tree = parse_state->tree_data->tree;
+  parse_state->tree_data->tree = new_tree;
+  system_mutex_release(parse_state->tree_data->tree_mutex);
+  return old_buffer_tree;
+}
+
+function void
+tree_sitter_parse_full_file_async(Async_Context* actx, String_Const_u8 data)
+{
+  if (data.size != sizeof(Buffer_ID)) return;
+  Buffer_ID buffer_id = *(Buffer_ID*)data.str;
+  Application_Links *app = actx->app;
+
+  acquire_global_frame_mutex(app);
+  Tree_Sitter_Parse_State parse_state = tree_sitter_parse_state_create(app, buffer_id);
+  release_global_frame_mutex(app);
+
   // 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);
+    new_tree = ts_parser_parse_string(parse_state.parser, parse_state.old_tree, (char *)parse_state.buffer_contents.str, (u32)parse_state.buffer_contents.size);
     if (async_check_canceled(actx))
     {
         canceled = true;
@@ -286,35 +312,35 @@ tree_sitter_parse_async__inner(Async_Context* actx, Buffer_ID buffer_id)
     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);
-
-      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);
+      old_buffer_tree = tree_sitter_parse_complete(&parse_state, new_tree);
+      animate_in_n_milliseconds(app, 0); // Force a frame refresh by requesting another frame
     }
     release_global_frame_mutex(app);
     ts_tree_delete(old_buffer_tree);
   }
-
-  ts_parser_delete(parser);
-  ts_tree_delete(old_tree);
-  linalloc_clear(&arena);
+  tree_sitter_parse_state_destroy(&parse_state);
 }
 
 function void
-tree_sitter_parse_async(Async_Context* actx, String_Const_u8 data)
+tree_sitter_parse_incremental(Application_Links* app, Buffer_ID buffer_id)
 {
-  if (data.size != sizeof(Buffer_ID)) return;
-  Buffer_ID buffer_id = *(Buffer_ID*)data.str;
-  tree_sitter_parse_async__inner(actx, buffer_id);
+  Tree_Sitter_Parse_State parse_state = tree_sitter_parse_state_create(app, buffer_id);
+  //Assert(parse_state.old_tree != 0); // if we are incrementally parsing, we expect there to be an old tree
+
+  TSTree *new_tree = 0;
+  while (!new_tree)
+  {
+    new_tree = ts_parser_parse_string(
+      parse_state.parser,
+      parse_state.old_tree,
+      (char *)parse_state.buffer_contents.str,
+      (u32)parse_state.buffer_contents.size
+    );
+  }
+
+  TSTree* old_buffer_tree = tree_sitter_parse_complete(&parse_state, new_tree);
+  ts_tree_delete(old_buffer_tree);
+  tree_sitter_parse_state_destroy(&parse_state);
 }
 
 ////////////////////////////////////////////////////////////////////
@@ -343,149 +369,205 @@ code_index_new_note(Code_Index_File* index, Arena* arena, Code_Index_Note_Kind k
   return result;
 }
 
-function Tree_Sitter_Code_Index_Parse_State
-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 = {};
-    parse_state.ok = false;
-    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);
-    parse_state.query = tree_sitter_query_init(
-      app, parse_state.buffer_id, parse_state.language->queries.ptr[Tree_Sitter_Language_Query_Tags]
+function Tree_Sitter_Code_Index_Update_State
+tree_sitter_code_index_update_state_create(
+  Application_Links* app,
+  Buffer_ID buffer_id,
+  Arena* scratch
+){
+  Tree_Sitter_Code_Index_Update_State state = {};
+  state.ok = false;
+  state.buffer_id = buffer_id;
+  state.language = tree_sitter_language_for_buffer(app, state.buffer_id);
+  if (!state.language) return state;
+  state.index_arena = make_arena_system(KB(16));
+  state.index = push_array_zero(&state.index_arena, Code_Index_File, 1);
+  state.buffer_contents = push_whole_buffer(app, scratch, state.buffer_id);
+  state.query = tree_sitter_query_init(
+    app, state.buffer_id, state.language->queries.ptr[Tree_Sitter_Language_Query_Tags]
   );
-    parse_state.nest_stack_first = 0;
-    parse_state.nest_stack_last = 0;
-    parse_state.last_note = 0;
-    parse_state.last_note_match_id = max_u32;
-  release_global_frame_mutex(app);
+  state.nest_stack_first = 0;
+  state.nest_stack_last = 0;
+  state.last_note = 0;
+  state.last_note_match_id = max_u32;
+  state.ok = true;
+  return state;
+}
 
-  parse_state.ok = true;
-  TSQueryMatch query_match;
-  u32 capture_index;
-  while (parse_state.ok && tree_sitter_query_continue(&parse_state.query, &query_match, &capture_index))
-  {
+function void
+tree_sitter_code_index_update_process_query_match(
+  Tree_Sitter_Code_Index_Update_State* state,
+  TSQueryMatch query_match,
+  u32 capture_index,
+  Arena* scratch
+){
   TSQueryCapture type_capture = query_match.captures[capture_index];
 
   TSNode type_node = type_capture.node;
   Range_i64 type_range = tree_sitter_node_to_range(type_node);
 
   u32 length;
-    const char* tmp = ts_query_capture_name_for_id(parse_state.query.query, type_capture.index, &length);
+  const char* tmp = ts_query_capture_name_for_id(state->query.query, type_capture.index, &length);
   String_Const_u8 capture_name = SCu8((char*)tmp, length);
 
   if (string_match(capture_name, SCu8("scope_begin")))
   {
-      Code_Index_Nest* nest = push_array_zero(&parse_state.index_arena, Code_Index_Nest, 1);
+    Code_Index_Nest* nest = push_array_zero(&state->index_arena, Code_Index_Nest, 1);
     nest->kind = CodeIndexNest_Scope;
     nest->open = type_range;
     nest->close = Ii64(max_i64);
-      nest->file = parse_state.index;
-      if (parse_state.nest_stack_last != 0)
+    nest->file = state->index;
+    if (state->nest_stack_last != 0)
     {
-        nest->parent = parse_state.nest_stack_last->nest;
+      nest->parent = state->nest_stack_last->nest;
       code_index_push_nest(&nest->parent->nest_list, nest);
     }
 
     Code_Index_Nest_Stack* stack = push_array_zero(scratch, Code_Index_Nest_Stack, 1);
     stack->nest = nest;
-      stack->prev = parse_state.nest_stack_last;
+    stack->prev = state->nest_stack_last;
     stack->match_id = query_match.id;
-      if (parse_state.nest_stack_last != 0) parse_state.nest_stack_last->next = stack;
-      else                      parse_state.nest_stack_first      = stack;
-      parse_state.nest_stack_last = stack;
+    if (state->nest_stack_last != 0) state->nest_stack_last->next = stack;
+    else                             state->nest_stack_first      = stack;
+    state->nest_stack_last = stack;
   }
   else if (string_match(capture_name, SCu8("scope_end")))
   {
-      Assert(parse_state.nest_stack_last != 0);
-      Assert(parse_state.nest_stack_last->match_id == query_match.id);
+    Assert(state->nest_stack_last != 0);
+    Assert(state->nest_stack_last->match_id == query_match.id);
 
-      Code_Index_Nest* nest = parse_state.nest_stack_last->nest;
+    Code_Index_Nest* nest = state->nest_stack_last->nest;
     nest->close = type_range;
     nest->is_closed = true;
-      nest->nest_array = code_index_nest_ptr_array_from_list(&parse_state.index_arena, &nest->nest_list);
-      if (nest->parent == 0) code_index_push_nest(&parse_state.index->nest_list, nest);
+    nest->nest_array = code_index_nest_ptr_array_from_list(&state->index_arena, &nest->nest_list);
+    if (nest->parent == 0) code_index_push_nest(&state->index->nest_list, nest);
 
-      parse_state.nest_stack_last = parse_state.nest_stack_last->prev;
-      if (parse_state.nest_stack_last != 0) parse_state.nest_stack_last->next = 0;
+    state->nest_stack_last = state->nest_stack_last->prev;
+    if (state->nest_stack_last != 0) state->nest_stack_last->next = 0;
   }
   else if (string_match(capture_name, SCu8("definition.class")))
   {
-      parse_state.last_note = code_index_new_note(parse_state.index, &parse_state.index_arena, CodeIndexNote_Type, type_range, parse_state.nest_stack_last);
-      parse_state.last_note_match_id = query_match.id;
+    state->last_note = code_index_new_note(state->index, &state->index_arena, CodeIndexNote_Type, type_range, state->nest_stack_last);
+    state->last_note_match_id = query_match.id;
   }
   else if (string_match(capture_name, SCu8("definition.function")))
   {
-      parse_state.last_note = code_index_new_note(parse_state.index, &parse_state.index_arena, CodeIndexNote_Function, type_range, parse_state.nest_stack_last);
-      parse_state.last_note_match_id = query_match.id;
+    state->last_note = code_index_new_note(state->index, &state->index_arena, CodeIndexNote_Function, type_range, state->nest_stack_last);
+    state->last_note_match_id = query_match.id;
   }
   else if (string_match(capture_name, SCu8("definition.method")))
   {
-      parse_state.last_note = code_index_new_note(parse_state.index, &parse_state.index_arena, CodeIndexNote_Function, type_range, parse_state.nest_stack_last);;
-      parse_state.last_note_match_id = query_match.id;
+    state->last_note = code_index_new_note(state->index, &state->index_arena, CodeIndexNote_Function, type_range, state->nest_stack_last);;
+    state->last_note_match_id = query_match.id;
   }
   else if (string_match(capture_name, SCu8("definition.type")))
   {
-      parse_state.last_note = code_index_new_note(parse_state.index, &parse_state.index_arena, CodeIndexNote_Type, type_range, parse_state.nest_stack_last);
-      parse_state.last_note_match_id = query_match.id;
+    state->last_note = code_index_new_note(state->index, &state->index_arena, CodeIndexNote_Type, type_range, state->nest_stack_last);
+    state->last_note_match_id = query_match.id;
   }
   else if (string_match(capture_name, SCu8("name")))
   {
-      if (parse_state.last_note != 0 && parse_state.last_note_match_id == query_match.id)
+    if (state->last_note != 0 && state->last_note_match_id == query_match.id)
     {
-        parse_state.last_note->pos = Ii64_size(type_range.start, type_range.end - type_range.start);
+      state->last_note->pos = Ii64_size(type_range.start, type_range.end - type_range.start);
+    }
   }
 }
 
+function void
+tree_sitter_code_index_update_complete(
+  Application_Links* app,
+  Tree_Sitter_Code_Index_Update_State* state
+){
+  tree_sitter_query_end(&state->query);
+  if (state->ok)
+  {
+    while (state->nest_stack_last != 0)
+    {
+      Code_Index_Nest* nest = state->nest_stack_last->nest;
+      if (nest->parent != 0) code_index_push_nest(&nest->parent->nest_list, nest);
+      else                   code_index_push_nest(&state->index->nest_list, nest);
+      nest->nest_array = code_index_nest_ptr_array_from_list(&state->index_arena, &nest->nest_list);
+      state->nest_stack_last = state->nest_stack_last->prev;
+    }
+
+    for (Code_Index_Note* note = state->index->note_list.first; note != 0 && note->next != note; note = note->next)
+    {
+      note->text = push_string_copy(&state->index_arena, string_substring(state->buffer_contents, note->pos));
+    }
+
+    // Finish the Index
+    state->index->nest_array = code_index_nest_ptr_array_from_list(&state->index_arena, &state->index->nest_list);
+    state->index->note_array = code_index_note_ptr_array_from_list(&state->index_arena, &state->index->note_list);
+
+
+    code_index_lock();
+    code_index_set_file(state->buffer_id, state->index_arena, state->index);
+    code_index_unlock();
+
+    buffer_clear_layout_cache(app, state->buffer_id);
+  }
+  else
+  {
+    linalloc_clear(&state->index_arena);
+  }
+}
+
+function Tree_Sitter_Code_Index_Update_State
+tree_sitter_code_index_update_full_file_async(Async_Context* actx, Buffer_ID buffer_id)
+{
+  Application_Links* app = actx->app;
+  Scratch_Block scratch(app);
+
+  acquire_global_frame_mutex(app);
+  Tree_Sitter_Code_Index_Update_State state = tree_sitter_code_index_update_state_create(
+    app,
+    buffer_id,
+    scratch
+  );
+  release_global_frame_mutex(app);
+
+  TSQueryMatch query_match;
+  u32 capture_index;
+  while (state.ok && tree_sitter_query_continue(&state.query, &query_match, &capture_index))
+  {
+    tree_sitter_code_index_update_process_query_match(&state, query_match, capture_index, scratch);
     if (async_check_canceled(actx))
     {
-      parse_state.ok = false;
+      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;
-      if (nest->parent != 0) code_index_push_nest(&nest->parent->nest_list, nest);
-      else                   code_index_push_nest(&parse_state.index->nest_list, nest);
-      nest->nest_array = code_index_nest_ptr_array_from_list(&parse_state.index_arena, &nest->nest_list);
-      parse_state.nest_stack_last = parse_state.nest_stack_last->prev;
+  tree_sitter_code_index_update_complete(app, &state);
+  return state;
 }
 
-    for (Code_Index_Note* note = parse_state.index->note_list.first; note != 0 && note->next != note; note = note->next)
+function Tree_Sitter_Code_Index_Update_State
+tree_sitter_code_index_update_full_file(Application_Links* app, Buffer_ID buffer_id)
 {
-      note->text = push_string_copy(&parse_state.index_arena, string_substring(parse_state.buffer_contents, note->pos));
-    }
+  Scratch_Block scratch(app);
 
-    // Finish the Index
-    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_Code_Index_Update_State state = tree_sitter_code_index_update_state_create(
+    app,
+    buffer_id,
+    scratch
+  );
 
-    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
+  TSQueryMatch query_match;
+  u32 capture_index;
+  while (state.ok && tree_sitter_query_continue(&state.query, &query_match, &capture_index))
   {
-    linalloc_clear(&parse_state.index_arena);
+    tree_sitter_code_index_update_process_query_match(
+      &state,
+      query_match,
+      capture_index,
+      scratch
+    );
   }
 
-  return parse_state;
+  tree_sitter_code_index_update_complete(app, &state);
+  return state;
 }
 
 function void
@@ -493,21 +575,23 @@ 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);
+  tree_sitter_code_index_update_full_file_async(actx, buffer_id);
 }
 
 function void
 tree_sitter_code_index_update_tick(Application_Links* app)
 {
+  Scratch_Block scratch(app);
   for (Buffer_Modified_Node* modified_node = global_buffer_modified_set.first;
        modified_node != 0;
        modified_node = modified_node->next
   ){
     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));
+    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);
+    if (!tree_data || !tree_data->tree) continue;
+
+    tree_sitter_code_index_update_full_file(app, buffer_id);
   }
   buffer_modified_set_clear();
 }

code/custom/4coder_tree_sitter.h

@@ -45,6 +45,10 @@ struct Buffer_Tree_Sitter_Data
 {
   TSTree* tree;
   System_Mutex tree_mutex;
+
+  // Code Index Update Requests
+  Range_i64 last_update_old_range;
+  Range_i64 last_update_new_range;
 };
 
 struct Tree_Sitter_Query_Cursor
@@ -65,7 +69,19 @@ struct Code_Index_Nest_Stack
   u32 match_id;
 };
 
-struct Tree_Sitter_Code_Index_Parse_State
+struct Tree_Sitter_Parse_State
+{
+  String_Const_u8 buffer_name;
+  Tree_Sitter_Language_Definition* lang;
+  String_Const_u8 buffer_contents;
+  Buffer_Tree_Sitter_Data* tree_data;
+
+  Arena arena;
+  TSParser* parser;
+  TSTree *old_tree;
+};
+
+struct Tree_Sitter_Code_Index_Update_State
 {
   Buffer_ID       buffer_id;
   String_Const_u8 buffer_contents;
@@ -104,7 +120,7 @@ function TSQuery* tree_sitter_query_new(Application_Links* app, TSLanguage* lang
 function void tree_sitter_register_language(String_Const_u8 ext, TSLanguage* language, Tree_Sitter_Language_Queries queries, Tree_Sitter_Language_Flags flags);
 function Tree_Sitter_Language_Definition* tree_sitter_language_for_buffer(Application_Links* app, Buffer_ID buffer_id);
 
-function Tree_Sitter_Code_Index_Parse_State tree_sitter_code_index_update_single_buffer(Application_Links *app, Buffer_ID buffer_id);
+function Tree_Sitter_Code_Index_Update_State tree_sitter_code_index_update_full_file(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);