Handle allocation failures when instantiating documents

src/runtime/alloc.h

@@ -1,12 +1,37 @@
 #ifndef RUNTIME_ALLOC_H_
 #define RUNTIME_ALLOC_H_
 
-#include <stdlib.h>
-
 #ifdef __cplusplus
 extern "C" {
 #endif
 
+#if defined(TREE_SITTER_WRAP_MALLOC)
+
+void *ts_record_malloc(size_t);
+void *ts_record_calloc(size_t, size_t);
+void *ts_record_realloc(void *, size_t);
+void ts_record_free(void *);
+
+static inline void *ts_malloc(size_t size) {
+  return ts_record_malloc(size);
+}
+
+static inline void *ts_calloc(size_t count, size_t size) {
+  return ts_record_calloc(count, size);
+}
+
+static inline void *ts_realloc(void *buffer, size_t size) {
+  return ts_record_realloc(buffer, size);
+}
+
+static inline void ts_free(void *buffer) {
+  return ts_record_free(buffer);
+}
+
+#else
+
+#include <stdlib.h>
+
 static inline void *ts_malloc(size_t size) {
   return malloc(size);
 }
@@ -23,6 +48,8 @@ static inline void ts_free(void *buffer) {
   return free(buffer);
 }
 
+#endif
+
 #ifdef __cplusplus
 }
 #endif

src/runtime/document.c

@@ -7,9 +7,16 @@
 #include "runtime/document.h"
 
 TSDocument *ts_document_make() {
-  TSDocument *document = ts_calloc(1, sizeof(TSDocument));
-  document->parser = ts_parser_make();
-  return document;
+  TSDocument *self = ts_calloc(1, sizeof(TSDocument));
+  if (!self)
+    return NULL;
+
+  if (!ts_parser_init(&self->parser)) {
+    ts_free(self);
+    return NULL;
+  }
+
+  return self;
 }
 
 void ts_document_free(TSDocument *self) {

src/runtime/parser.c

@@ -637,18 +637,35 @@ static bool ts_parser__consume_lookahead(TSParser *self, int head,
  *  Public
  */
 
-TSParser ts_parser_make() {
-  return (TSParser){
-    .lexer = ts_lexer_make(),
-    .stack = ts_stack_new(),
-    .lookahead_states = vector_new(sizeof(LookaheadState), 4),
-    .reduce_parents = vector_new(sizeof(TSTree *), 4),
-    .finished_tree = NULL,
-  };
+bool ts_parser_init(TSParser *self) {
+  self->finished_tree = NULL;
+  self->lexer = ts_lexer_make();
+
+  self->stack = ts_stack_new();
+  if (!self->stack) {
+    return false;
+  }
+
+  self->lookahead_states = vector_new(sizeof(LookaheadState), 4);
+  if (!self->lookahead_states.contents) {
+    ts_stack_delete(self->stack);
+    return false;
+  }
+
+  self->reduce_parents = vector_new(sizeof(TSTree *), 4);
+  if (!self->reduce_parents.contents) {
+    ts_stack_delete(self->stack);
+    vector_delete(&self->lookahead_states);
+    return false;
+  }
+
+  return true;
 }
 
 void ts_parser_destroy(TSParser *self) {
   ts_stack_delete(self->stack);
+  vector_delete(&self->lookahead_states);
+  vector_delete(&self->reduce_parents);
 }
 
 TSDebugger ts_parser_debugger(const TSParser *self) {

src/runtime/parser.h

@@ -18,7 +18,7 @@ typedef struct {
   bool is_split;
 } TSParser;
 
-TSParser ts_parser_make();
+bool ts_parser_init(TSParser *);
 void ts_parser_destroy(TSParser *);
 TSDebugger ts_parser_debugger(const TSParser *);
 void ts_parser_set_debugger(TSParser *, TSDebugger);

src/runtime/stack.c

@@ -43,17 +43,39 @@ static TSTree *ts_stack__default_tree_selection(void *p, TSTree *t1, TSTree *t2)
 }
 
 Stack *ts_stack_new() {
-  Stack *self = ts_malloc(sizeof(Stack));
-  *self = (Stack){
-    .heads = ts_calloc(INITIAL_HEAD_CAPACITY, sizeof(StackNode *)),
-    .head_count = 1,
-    .head_capacity = INITIAL_HEAD_CAPACITY,
-    .tree_selection_payload = NULL,
-    .tree_selection_function = ts_stack__default_tree_selection,
-    .pop_results = vector_new(sizeof(StackPopResult), 4),
-    .pop_paths = vector_new(sizeof(PopPath), 4),
-  };
+  Stack *self = ts_calloc(1, sizeof(Stack));
+  if (!self)
+    goto error;
+
+  self->head_count = 1;
+  self->head_capacity = INITIAL_HEAD_CAPACITY;
+  self->heads = ts_calloc(INITIAL_HEAD_CAPACITY, sizeof(StackNode *));
+  if (!self->heads)
+    goto error;
+
+  self->pop_results = vector_new(sizeof(StackPopResult), 4);
+  if (!self->pop_results.contents)
+    goto error;
+
+  self->pop_paths = vector_new(sizeof(PopPath), 4);
+  if (!self->pop_paths.contents)
+    goto error;
+
+  self->tree_selection_payload = NULL;
+  self->tree_selection_function = ts_stack__default_tree_selection;
   return self;
+
+error:
+  if (self) {
+    if (self->heads)
+      ts_free(self->heads);
+    if (self->pop_results.contents)
+      vector_delete(&self->pop_results);
+    if (self->pop_paths.contents)
+      vector_delete(&self->pop_paths);
+    ts_free(self);
+  }
+  return NULL;
 }
 
 void ts_stack_delete(Stack *self) {