treesitter_elixir

#include "erl_nif.h"
#include "tree_sitter/api.h"
#include <dlfcn.h>
#include <stdint.h>

// Setup external resources
static ErlNifResourceType *parser_resource_type = NULL;
static ErlNifResourceType *tree_resource_type = NULL;
static ErlNifResourceType *node_resource_type = NULL;
static ErlNifResourceType *language_resource_type = NULL;
static ErlNifResourceType *query_resource_type = NULL;
static ErlNifResourceType *querycursor_resource_type = NULL;

// Wrappers
typedef struct {
  TSParser *data;
} WRParser;

typedef struct {
  TSTree *data;
} WRTree;

typedef struct {
  // NOTE:TSNode points to the tree, so there could be a risk of the tree being
  // freed first. We could have a wrapper that has the parser, the tree, and the
  // current node?
  // Could also hold the source_code/text
  // I think this is why the cursor stuff exists, so maybe rely on that.
  TSNode data;
} WRNode;

typedef struct {
  TSLanguage *data;
} WRLanguage;

typedef struct {
  TSQuery *data;
} WRQuery;

typedef struct {
  TSQueryCursor *data;
} WRQueryCursor;

static void parser_dtor(ErlNifEnv *env, void *obj) {
  WRParser *my_struct = (WRParser *)obj;
  if (my_struct->data) {
    ts_parser_delete(my_struct->data);
  }
  enif_release_resource(my_struct);
}

static void tree_dtor(ErlNifEnv *env, void *obj) {
  WRTree *my_struct = (WRTree *)obj;
  if (my_struct->data) {
    ts_tree_delete(my_struct->data);
  }
  enif_release_resource(my_struct);
}

static void node_dtor(ErlNifEnv *env, void *obj) {
  WRNode *my_struct = (WRNode *)obj;
  // NOTE: Check with mistral if this is seems correct
  enif_release_resource(my_struct);
}

static void language_dtor(ErlNifEnv *env, void *obj) {
  WRLanguage *my_struct = (WRLanguage *)obj;
  if (my_struct->data) {
    ts_language_delete(my_struct->data);
  }
  enif_release_resource(my_struct);
}

static void query_dtor(ErlNifEnv *env, void *obj) {
  WRQuery *my_struct = (WRQuery *)obj;
  if (my_struct->data) {
    ts_query_delete(my_struct->data);
  }
  enif_release_resource(my_struct);
}

static void querycursor_dtor(ErlNifEnv *env, void *obj) {
  WRQueryCursor *my_struct = (WRQueryCursor *)obj;
  if (my_struct->data) {
    ts_query_cursor_delete(my_struct->data);
  }
  enif_release_resource(my_struct);
}

static int load(ErlNifEnv *env, void **priv_data, ERL_NIF_TERM load_info) {
  ErlNifResourceFlags flags =
      (ErlNifResourceFlags)(ERL_NIF_RT_CREATE | ERL_NIF_RT_TAKEOVER);
  parser_resource_type = enif_open_resource_type(env, NULL, "parser_resource",
                                                 parser_dtor, flags, NULL);

  tree_resource_type = enif_open_resource_type(env, NULL, "tree_resource",
                                               tree_dtor, flags, NULL);

  node_resource_type = enif_open_resource_type(env, NULL, "node_resource",
                                               node_dtor, flags, NULL);
  language_resource_type = enif_open_resource_type(
      env, NULL, "language_resource", language_dtor, flags, NULL);
  query_resource_type = enif_open_resource_type(env, NULL, "query_resource",
                                                query_dtor, flags, NULL);
  querycursor_resource_type = enif_open_resource_type(
      env, NULL, "querycursor_resource", querycursor_dtor, flags, NULL);
  return 0;
}

ERL_NIF_TERM to_error_string(ErlNifEnv *env, const char *msg) {
  return enif_make_tuple2(env, enif_make_atom(env, "error"),
                          enif_make_string(env, msg, ERL_NIF_UTF8));
}

ERL_NIF_TERM to_error_nil(ErlNifEnv *env) {
  return enif_make_tuple2(env, enif_make_atom(env, "error"),
                          enif_make_atom(env, "nil"));
}

ERL_NIF_TERM to_ok_term(ErlNifEnv *env, ERL_NIF_TERM value) {
  return enif_make_tuple2(env, enif_make_atom(env, "ok"), value);
}

static ERL_NIF_TERM parser_new(ErlNifEnv *env, int argc,
                               const ERL_NIF_TERM argv[]) {
  WRParser *wrap = enif_alloc_resource(parser_resource_type, sizeof(WRParser));
  TSParser *parser = ts_parser_new();
  if (!parser) {
    return enif_make_atom(env, "error");
  }
  wrap->data = parser;
  return to_ok_term(env, enif_make_resource(env, wrap));
  // return enif_make_atom(env, "error");
}

typedef TSLanguage *(tree_sitter_lang)(void);
static ERL_NIF_TERM language(ErlNifEnv *env, int argc,
                             const ERL_NIF_TERM argv[]) {
  char language[1024];
  char lib_path[1024];
  if (!enif_get_string(env, argv[0], language, 1024, ERL_NIF_UTF8)) {
    return to_error_string(env, language);
  }
  if (!enif_get_string(env, argv[1], lib_path, 1024, ERL_NIF_UTF8)) {
    return to_error_string(env, lib_path);
  }

  void *lib = dlopen(lib_path, RTLD_NOW);
  const char *err = dlerror();
  if (err != NULL) {
    return to_error_string(env, "Unable to load library");
  }

  tree_sitter_lang *make_ts_language = (tree_sitter_lang *)dlsym(lib, language);
  WRLanguage *wrap =
      enif_alloc_resource(language_resource_type, sizeof(WRLanguage));
  wrap->data = make_ts_language();
  return to_ok_term(env, enif_make_resource(env, wrap));
}

static ERL_NIF_TERM parser_set_language(ErlNifEnv *env, int argc,
                                        const ERL_NIF_TERM argv[]) {
  WRParser *wrap;
  if (!enif_get_resource(env, argv[0], parser_resource_type, (void **)&wrap)) {
    return to_error_string(env, "Unable to unwrap parser");
  }

  WRLanguage *lang_wrap;
  if (!enif_get_resource(env, argv[1], language_resource_type,
                         (void **)&lang_wrap)) {
    return to_error_string(env, "Unable to unwrap language");
  }

  if (!ts_parser_set_language(wrap->data, lang_wrap->data)) {
    return to_error_string(env, "Unable to set the language to the parser");
  }

  return to_ok_term(env, enif_make_resource(env, wrap));
}

static ERL_NIF_TERM parse(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
  WRParser *wrap;
  if (!enif_get_resource(env, argv[0], parser_resource_type, (void **)&wrap)) {
    return to_error_string(env, "Unable to unwrap parser");
  }

  unsigned int len;
  enif_get_string_length(env, argv[1], &len, ERL_NIF_UTF8);
  char *source_code = malloc(len + 1);
  if (!enif_get_string(env, argv[1], source_code, len + 1, ERL_NIF_UTF8)) {
    return to_error_string(env, "Failed loading source code");
  }

  TSTree *tree = ts_parser_parse_string(wrap->data, NULL, source_code,
                                        strlen(source_code));
  if (!tree) {
    return to_error_string(env, "Failed parsing the source code");
  }
  WRTree *wraptree = enif_alloc_resource(tree_resource_type, sizeof(WRTree));
  wraptree->data = tree;
  free(source_code);
  return to_ok_term(env, enif_make_resource(env, wraptree));
}
static ERL_NIF_TERM enif_make_node_wrap(ErlNifEnv *env, const TSNode node) {
  WRNode *wrapnode = enif_alloc_resource(node_resource_type, sizeof(WRNode));
  wrapnode->data = node;
  return enif_make_resource(env, wrapnode);
}

static ERL_NIF_TERM tree_root_node(ErlNifEnv *env, int argc,
                                   const ERL_NIF_TERM argv[]) {
  WRTree *wrap;
  if (!enif_get_resource(env, argv[0], tree_resource_type, (void **)&wrap)) {
    return to_error_string(env, "Unable to unwrap tree");
  }
  TSNode root_node = ts_tree_root_node(wrap->data);
  WRNode *wrapnode = enif_alloc_resource(node_resource_type, sizeof(WRNode));
  wrapnode->data = root_node;
  return to_ok_term(env, enif_make_resource(env, wrapnode));
}

static ERL_NIF_TERM node_parent(ErlNifEnv *env, int argc,
                                const ERL_NIF_TERM argv[]) {
  WRNode *wrap;
  if (!enif_get_resource(env, argv[0], node_resource_type, (void **)&wrap)) {
    return to_error_string(env, "Unable to unwrap node");
  }
  TSNode node = ts_node_parent(wrap->data);
  if (ts_node_is_null(node)) {
    return to_error_nil(env);
  } else {
    WRNode *wrapnode = enif_alloc_resource(node_resource_type, sizeof(WRNode));
    wrapnode->data = node;
    return to_ok_term(env, enif_make_resource(env, wrapnode));
  }
}

static ERL_NIF_TERM node_child(ErlNifEnv *env, int argc,
                               const ERL_NIF_TERM argv[]) {
  WRNode *wrap;
  if (!enif_get_resource(env, argv[0], node_resource_type, (void **)&wrap)) {
    return to_error_string(env, "Unable to unwrap node");
  }
  TSNode node = ts_node_child(wrap->data, 0);
  if (ts_node_is_null(node)) {
    return to_error_nil(env);
  } else {
    WRNode *wrapnode = enif_alloc_resource(node_resource_type, sizeof(WRNode));
    wrapnode->data = node;
    return to_ok_term(env, enif_make_resource(env, wrapnode));
  }
}

static ERL_NIF_TERM node_next_sibling(ErlNifEnv *env, int argc,
                                      const ERL_NIF_TERM argv[]) {
  WRNode *wrap;
  if (!enif_get_resource(env, argv[0], node_resource_type, (void **)&wrap)) {
    return to_error_string(env, "Unable to unwrap node");
  }
  TSNode node = ts_node_next_sibling(wrap->data);
  if (ts_node_is_null(node)) {
    return to_error_nil(env);
  } else {
    WRNode *wrapnode = enif_alloc_resource(node_resource_type, sizeof(WRNode));
    wrapnode->data = node;
    return to_ok_term(env, enif_make_resource(env, wrapnode));
  }
}

static ERL_NIF_TERM node_prev_sibling(ErlNifEnv *env, int argc,
                                      const ERL_NIF_TERM argv[]) {
  WRNode *wrap;
  if (!enif_get_resource(env, argv[0], node_resource_type, (void **)&wrap)) {
    return to_error_string(env, "Unable to unwrap node");
  }
  TSNode node = ts_node_prev_sibling(wrap->data);
  if (ts_node_is_null(node)) {
    return to_error_nil(env);
  } else {
    WRNode *wrapnode = enif_alloc_resource(node_resource_type, sizeof(WRNode));
    wrapnode->data = node;
    return to_ok_term(env, enif_make_resource(env, wrapnode));
  }
}

static ERL_NIF_TERM node_string(ErlNifEnv *env, int argc,
                                const ERL_NIF_TERM argv[]) {
  WRNode *wrap;
  if (!enif_get_resource(env, argv[0], node_resource_type, (void **)&wrap)) {
    return to_error_string(env, "Unable to unwrap tree");
  }
  char *string = ts_node_string(wrap->data);

  // Free all of the heap-allocated memory.
  ERL_NIF_TERM type = enif_make_string(env, string, ERL_NIF_UTF8);
  free(string);
  return type;
}

static ERL_NIF_TERM node_type(ErlNifEnv *env, int argc,
                              const ERL_NIF_TERM argv[]) {
  WRNode *wrap;
  if (!enif_get_resource(env, argv[0], node_resource_type, (void **)&wrap)) {
    return to_error_string(env, "Unable to unwrap tree");
  }
  const char *string = ts_node_type(wrap->data);

  // Free all of the heap-allocated memory.
  ERL_NIF_TERM type = enif_make_string(env, string, ERL_NIF_UTF8);
  return type;
}

static ERL_NIF_TERM node_byte_range(ErlNifEnv *env, int argc,
                                    const ERL_NIF_TERM argv[]) {
  WRNode *wrap;
  if (!enif_get_resource(env, argv[0], node_resource_type, (void **)&wrap)) {
    return to_error_string(env, "Unable to unwrap tree");
  }
  int i = ts_node_start_byte(wrap->data);
  int j = ts_node_end_byte(wrap->data);
  return enif_make_tuple2(env, enif_make_int(env, i), enif_make_int(env, j));
}

ERL_NIF_TERM point_to_tuple(ErlNifEnv *env, TSPoint point) {
  ERL_NIF_TERM map = enif_make_new_map(env);

  // Add the 'row' key-value pair to the map
  ERL_NIF_TERM row_key = enif_make_atom(env, "row");
  ERL_NIF_TERM row_value = enif_make_int(env, point.row);
  enif_make_map_put(env, map, row_key, row_value, &map);

  // Add the 'column' key-value pair to the map
  ERL_NIF_TERM column_key = enif_make_atom(env, "column");
  ERL_NIF_TERM column_value = enif_make_int(env, point.column);
  enif_make_map_put(env, map, column_key, column_value, &map);
  return map;
}

static ERL_NIF_TERM node_point_range(ErlNifEnv *env, int argc,
                                     const ERL_NIF_TERM argv[]) {
  WRNode *wrap;
  if (!enif_get_resource(env, argv[0], node_resource_type, (void **)&wrap)) {
    return to_error_string(env, "Unable to unwrap tree");
  }
  TSPoint i = ts_node_start_point(wrap->data);
  TSPoint j = ts_node_end_point(wrap->data);
  return enif_make_tuple2(env, point_to_tuple(env, i), point_to_tuple(env, j));
}

/**
 * Query
 */
// TODO: Have a tslanguage wrapper and support getting it
// integrate this above
// TSQueryCapture to ERL_NIF_TERM (list of nodes)
// TSQueryMatch to ERL_NIF_TERM (dict with pattern_index and the list of
// nodes(tsqueraycapture)

// TSQuery *ts_query_new(
//   const TSLanguage *language,
//   const char *source,
//   uint32_t source_len,
//   uint32_t *error_offset,
//   TSQueryError *error_type
// );

static ERL_NIF_TERM query_new(ErlNifEnv *env, int argc,
                              const ERL_NIF_TERM argv[]) {
  WRLanguage *lang_wrap;
  if (!enif_get_resource(env, argv[0], language_resource_type,
                         (void **)&lang_wrap)) {
    return to_error_string(env, "Unable to unwrap language");
  }
  unsigned int len;
  enif_get_string_length(env, argv[1], &len, ERL_NIF_UTF8);
  char *source_code = malloc(len + 1);
  if (!enif_get_string(env, argv[1], source_code, len + 1, ERL_NIF_UTF8)) {
    return to_error_string(env, "Failed loading source code");
  }

  WRQuery *wrap = enif_alloc_resource(query_resource_type, sizeof(WRQuery));
  uint32_t error_offset;
  TSQueryError error_type;
  wrap->data = ts_query_new(lang_wrap->data, source_code, len, &error_offset,
                            &error_type);
  if (wrap->data) {
    return to_ok_term(env, enif_make_resource(env, wrap));
  }
  char error_string[50];
  sprintf(error_string, "Found an error in the query at %d, type %d",
          error_offset, error_type);

  return to_error_string(env, error_string);
}

// void ts_query_cursor_exec(TSQueryCursor *, const TSQuery *, TSNode);
static ERL_NIF_TERM query_cursor_exec(ErlNifEnv *env, int argc,
                                      const ERL_NIF_TERM argv[]) {
  WRQuery *query_wrap;
  if (!enif_get_resource(env, argv[0], query_resource_type,
                         (void **)&query_wrap)) {
    return to_error_string(env, "Unable to unwrap query");
  }
  WRNode *node_wrap;
  if (!enif_get_resource(env, argv[1], node_resource_type,
                         (void **)&node_wrap)) {
    return to_error_string(env, "Unable to unwrap node");
  }
  WRQueryCursor *wrap =
      enif_alloc_resource(querycursor_resource_type, sizeof(WRQueryCursor));
  wrap->data = ts_query_cursor_new();
  if (wrap->data) {
    ts_query_cursor_exec(wrap->data, query_wrap->data, node_wrap->data);

    return to_ok_term(env, enif_make_resource(env, wrap));
  }
  return to_error_string(env, "Unable to create cursor");
}

// bool ts_query_cursor_next_match(TSQueryCursor *, TSQueryMatch *match);
// TODO: Return the cursor?
static ERL_NIF_TERM query_cursor_next_match(ErlNifEnv *env, int argc,
                                            const ERL_NIF_TERM argv[]) {
  WRQueryCursor *querycursor_wrap;
  if (!enif_get_resource(env, argv[0], querycursor_resource_type,
                         (void **)&querycursor_wrap)) {
    return to_error_string(env, "Unable to unwrap querycursor");
  }
  TSQueryMatch match;
  if (!ts_query_cursor_next_match(querycursor_wrap->data, &match)) {
    return to_error_nil(env);
  }
  // Add them to an list
  ERL_NIF_TERM *lst =
      (ERL_NIF_TERM *)malloc(match.capture_count * sizeof(ERL_NIF_TERM));
  for (int i = 0; i < match.capture_count; ++i) {
    lst[i] = enif_make_node_wrap(env, match.captures[i].node);
  }
  ERL_NIF_TERM arr = enif_make_list_from_array(env, lst, match.capture_count);
  free(lst);
  // TODO: Do we ever need the match id and pattern_index?
  return enif_make_tuple2(
      env, enif_make_atom(env, "ok"),
      enif_make_tuple3(env, enif_make_int(env, match.id),
                       enif_make_int(env, match.pattern_index), arr));
}

static ERL_NIF_TERM query_pattern_byte_range(ErlNifEnv *env, int argc,
                                             const ERL_NIF_TERM argv[]) {
  // TODO: Can we make this into a function
  WRQuery *query_wrap;
  if (!enif_get_resource(env, argv[0], query_resource_type,
                         (void **)&query_wrap)) {
    return to_error_string(env, "Unable to unwrap query");
  }

  uint32_t pattern_index;
  enif_get_uint(env, argv[1], &pattern_index);

  uint32_t byte_start =
      ts_query_start_byte_for_pattern(query_wrap->data, pattern_index);
  uint32_t byte_end =
      ts_query_end_byte_for_pattern(query_wrap->data, pattern_index);
  return enif_make_tuple2(env, enif_make_int(env, byte_start),
                          enif_make_int(env, byte_end));
}

/**
 * Get the name and length of one of the query's captures, or one of the
 * query's string literals. Each capture and string is associated with a
 * numeric id based on the order that it appeared in the query's source.
 */
// const char *ts_query_capture_name_for_id(
//   const TSQuery *self,
//   uint32_t index,
//   uint32_t *length
// );
static ERL_NIF_TERM query_capture_name_for_id(ErlNifEnv *env, int argc,
                                              const ERL_NIF_TERM argv[]) {
  // TODO: Can we make this into a function
  WRQuery *query_wrap;
  if (!enif_get_resource(env, argv[0], query_resource_type,
                         (void **)&query_wrap)) {
    return to_error_string(env, "Unable to unwrap query");
  }

  uint32_t pattern_index;
  enif_get_uint(env, argv[1], &pattern_index);
  int len;
  const char *name =
      ts_query_capture_name_for_id(query_wrap->data, pattern_index, &len);
  return enif_make_tuple2(env, enif_make_atom(env, "ok"),
                          enif_make_string(env, name, ERL_NIF_UTF8));
}

static ErlNifFunc nif_funcs[] = {
    {"parser_new", 0, parser_new, 0},
    {"language", 2, language, 0},
    {"parser_set_language", 2, parser_set_language, 0},
    {"parse", 2, parse, 0},
    {"tree_root_node", 1, tree_root_node, 0},
    {"node_parent", 1, node_parent, 0},
    {"node_child", 1, node_child, 0},
    {"node_next_sibling", 1, node_next_sibling, 0},
    {"node_prev_sibling", 1, node_prev_sibling, 0},
    {"node_string", 1, node_string, 0},
    {"node_type", 1, node_type, 0},
    {"node_byte_range", 1, node_byte_range, 0},
    {"node_point_range", 1, node_point_range, 0},
    {"query_new", 2, query_new, 0},
    {"query_cursor_exec", 2, query_cursor_exec, 0},
    {"query_cursor_next_match", 1, query_cursor_next_match, 0},
    {"query_pattern_byte_range", 2, query_pattern_byte_range, 0},
    {"query_capture_name_for_id", 2, query_capture_name_for_id, 0}};

ERL_NIF_INIT(Elixir.TreeSitter.Nif, nif_funcs, load, NULL, NULL, NULL)