treesitter_elixir

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

// Setup external resources
static ErlNifResourceType *parser_resource_type = NULL;
static ErlNifResourceType *tree_resource_type = NULL;
static ErlNifResourceType *node_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
  TSNode data;
} WRNode;

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 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);
  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 parser_set_language(ErlNifEnv *env, int argc,
                                        const ERL_NIF_TERM argv[]) {
  char language[1024];
  char lib_path[1024];
  if (!enif_get_string(env, argv[1], language, 1024, ERL_NIF_UTF8)) {
    return to_error_string(env, language);
  }
  if (!enif_get_string(env, argv[2], lib_path, 1024, ERL_NIF_UTF8)) {
    return to_error_string(env, lib_path);
  }

  WRParser *wrap;
  if (!enif_get_resource(env, argv[0], parser_resource_type, (void **)&wrap)) {
    return to_error_string(env, "Unable to unwrap parser");
  }

  // Close lib properly, currently seems to result in segmentation vault.
  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);
  TSLanguage *lang = make_ts_language();
  if (!ts_parser_set_language(wrap->data, lang)) {
    dlclose(lib);
    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 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_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));
}

static ErlNifFunc nif_funcs[] = {
    {"parser_new", 0, parser_new, 0},
    {"parser_set_language", 3, 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_range", 1, node_range, 0},
    {"node_point_range", 1, node_point_range, 0}};

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