4coder/code/custom/4coder_tree_sitter.cpp

function bool
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);
}

function bool
tree_sitter_begin_buffer(Application_Links* app, Buffer_ID buffer_id, File_Language_Kind kind)
{
  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);
  switch (kind)
  {
    case File_Language_CPP:
    {
      tree_data->language = tree_sitter_cpp();
    } break;

    default:
      tree_data->language = 0;
  }

  if (tree_data->language != 0)
  {
    tree_data->tree_mutex = system_mutex_make();
  }

  return tree_data->language != 0;
}

function void
tree_sitter_end_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);
  if (!tree_data || !tree_data->language) return;

  Async_Task *tree_sitter_parse_task = scope_attachment(app, buffer_scope, buffer_tree_sitter_parse_task_id, Async_Task);
  if (async_task_is_running_or_pending(&global_async_system, *tree_sitter_parse_task))
  {
    async_task_cancel(app, &global_async_system, *tree_sitter_parse_task);
  }

  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);
    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, tree_data->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 void
tree_sitter_code_index_update_tick(Application_Links* app)
{

}

////////////////////////////////////////////////////////////////////
// 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);
  }
}