Allow wasm languages to be deleted

lib/src/language.c

@@ -1,13 +1,19 @@
 #include "./language.h"
+#include "./wasm.h"
+#include "tree_sitter/api.h"
 #include <string.h>
 
 const TSLanguage *ts_language_copy(const TSLanguage *self) {
-  // TODO - increment reference count for wasm languages
+  if (self && ts_language_is_wasm(self)) {
+    ts_wasm_language_retain(self);
+  }
   return self;
 }
 
-void ts_language_delete(const TSLanguage *_self) {
-  // TODO - decrement reference count for wasm languages
+void ts_language_delete(const TSLanguage *self) {
+  if (self && ts_language_is_wasm(self)) {
+    ts_wasm_language_release(self);
+  }
 }
 
 uint32_t ts_language_symbol_count(const TSLanguage *self) {

lib/src/parser.c

@@ -1869,6 +1869,7 @@ const TSLanguage *ts_parser_language(const TSParser *self) {
 
 bool ts_parser_set_language(TSParser *self, const TSLanguage *language) {
   ts_parser__external_scanner_destroy(self);
+  ts_language_delete(self->language);
   self->language = NULL;
 
   if (language) {
@@ -1885,7 +1886,7 @@ bool ts_parser_set_language(TSParser *self, const TSLanguage *language) {
     }
   }
 
-  self->language = language;
+  self->language = ts_language_copy(language);
   ts_parser__external_scanner_create(self);
   ts_parser_reset(self);
   return true;

lib/src/query.c

@@ -2812,6 +2812,7 @@ void ts_query_delete(TSQuery *self) {
     array_delete(&self->string_buffer);
     array_delete(&self->negated_fields);
     array_delete(&self->repeat_symbols_with_rootless_patterns);
+    ts_language_delete(self->language);
     symbol_table_delete(&self->captures);
     symbol_table_delete(&self->predicate_values);
     for (uint32_t index = 0; index < self->capture_quantifiers.size; index++) {

lib/src/tree.c

@@ -12,7 +12,7 @@ TSTree *ts_tree_new(
 ) {
   TSTree *result = ts_malloc(sizeof(TSTree));
   result->root = root;
-  result->language = language;
+  result->language = ts_language_copy(language);
   result->included_ranges = ts_calloc(included_range_count, sizeof(TSRange));
   memcpy(result->included_ranges, included_ranges, included_range_count * sizeof(TSRange));
   result->included_range_count = included_range_count;
@@ -30,6 +30,7 @@ void ts_tree_delete(TSTree *self) {
   SubtreePool pool = ts_subtree_pool_new(0);
   ts_subtree_release(&pool, self->root);
   ts_subtree_pool_delete(&pool);
+  ts_language_delete(self->language);
   ts_free(self->included_ranges);
   ts_free(self);
 }

lib/src/wasm.c

@@ -72,6 +72,7 @@ typedef struct {
 // wasm store, so it can be shared by all users of a `TSLanguage`. A pointer to
 // this is stored on the language itself.
 typedef struct {
+  volatile uint32_t ref_count;
   wasmtime_module_t *module;
   uint32_t language_id;
   const char *name;
@@ -1074,7 +1075,7 @@ const TSLanguage *ts_wasm_store_load_language(
   };
   uint32_t address_count = array_len(addresses);
 
-  TSLanguage *language = ts_malloc(sizeof(TSLanguage));
+  TSLanguage *language = ts_calloc(1, sizeof(TSLanguage));
   StringData symbol_name_buffer = array_new();
   StringData field_name_buffer = array_new();
 
@@ -1202,6 +1203,7 @@ const TSLanguage *ts_wasm_store_load_language(
     .symbol_name_buffer = symbol_name_buffer.contents,
     .field_name_buffer = field_name_buffer.contents,
     .dylink_info = dylink_info,
+    .ref_count = 1,
   };
 
   // The lex functions are not used for wasm languages. Use those two fields
@@ -1468,6 +1470,45 @@ bool ts_language_is_wasm(const TSLanguage *self) {
   return self->lex_fn == ts_wasm_store__sentinel_lex_fn;
 }
 
+static inline LanguageWasmModule *ts_language__wasm_module(const TSLanguage *self) {
+  return (LanguageWasmModule *)self->keyword_lex_fn;
+}
+
+void ts_wasm_language_retain(const TSLanguage *self) {
+  LanguageWasmModule *module = ts_language__wasm_module(self);
+  assert(module->ref_count > 0);
+  atomic_inc(&module->ref_count);
+}
+
+void ts_wasm_language_release(const TSLanguage *self) {
+  LanguageWasmModule *module = ts_language__wasm_module(self);
+  assert(module->ref_count > 0);
+  if (atomic_dec(&module->ref_count) == 0) {
+    ts_free((void *)module->field_name_buffer);
+    ts_free((void *)module->symbol_name_buffer);
+    ts_free((void *)module->name);
+    wasmtime_module_delete(module->module);
+    ts_free(module);
+
+    ts_free((void *)self->alias_map);
+    ts_free((void *)self->alias_sequences);
+    ts_free((void *)self->external_scanner.symbol_map);
+    ts_free((void *)self->field_map_entries);
+    ts_free((void *)self->field_map_slices);
+    ts_free((void *)self->field_names);
+    ts_free((void *)self->lex_modes);
+    ts_free((void *)self->parse_actions);
+    ts_free((void *)self->parse_table);
+    ts_free((void *)self->primary_state_ids);
+    ts_free((void *)self->public_symbol_map);
+    ts_free((void *)self->small_parse_table);
+    ts_free((void *)self->small_parse_table_map);
+    ts_free((void *)self->symbol_metadata);
+    ts_free((void *)self->symbol_names);
+    ts_free((void *)self);
+  }
+}
+
 #else
 
 // If the WASM feature is not enabled, define dummy versions of all of the
@@ -1556,4 +1597,12 @@ bool ts_language_is_wasm(const TSLanguage *self) {
   return false;
 }
 
+void ts_wasm_language_retain(const TSLanguage *self) {
+  (void)self;
+}
+
+void ts_wasm_language_release(const TSLanguage *self) {
+  (void)self;
+}
+
 #endif

lib/src/wasm.h

@@ -20,6 +20,9 @@ bool ts_wasm_store_call_scanner_scan(TSWasmStore *, uint32_t, uint32_t);
 uint32_t ts_wasm_store_call_scanner_serialize(TSWasmStore *, uint32_t, char *);
 void ts_wasm_store_call_scanner_deserialize(TSWasmStore *, uint32_t, const char *, unsigned);
 
+void ts_wasm_language_retain(const TSLanguage *);
+void ts_wasm_language_release(const TSLanguage *);
+
 #ifdef __cplusplus
 }
 #endif