Handle allocation failures when instantiating documents

project.gyp

@@ -142,6 +142,15 @@
           'GCC_OPTIMIZATION_LEVEL': '0',
         },
       },
+      'Test': {
+        'defines': ['TREE_SITTER_WRAP_MALLOC=true'],
+        'cflags': [ '-g' ],
+        'ldflags': [ '-g' ],
+        'xcode_settings': {
+          'OTHER_LDFLAGS': ['-g'],
+          'GCC_OPTIMIZATION_LEVEL': '0',
+        },
+      },
       'Release': {
         'cflags': [ '-O2', '-fno-strict-aliasing' ],
         'cflags!': [ '-O3', '-fstrict-aliasing' ],

script/test

@@ -27,7 +27,8 @@ profile=
 mode=normal
 args=()
 target=tests
-cmd="out/Debug/${target}"
+export BUILDTYPE=Test
+cmd="out/${BUILDTYPE}/${target}"
 
 while getopts "df:s:ghpv" option; do
   case ${option} in
@@ -56,7 +57,7 @@ while getopts "df:s:ghpv" option; do
   esac
 done
 
-BUILDTYPE=Debug make $target
+make $target
 args=${args:-""}
 
 if [[ -n $profile ]]; then

spec/helpers/record_alloc.cc

@@ -0,0 +1,95 @@
+#include <stdlib.h>
+#include <map>
+#include <set>
+#include "bandit/bandit.h"
+
+using std::map;
+using std::set;
+
+bool _enabled = false;
+static size_t _allocation_count = 0;
+static map<void *, size_t> _outstanding_allocations;
+static size_t _allocation_failure_index = -1;
+
+namespace record_alloc {
+
+void start() {
+  _enabled = true;
+  _allocation_count = 0;
+  _outstanding_allocations.clear();
+  _allocation_failure_index = -1;
+}
+
+void stop() {
+  _enabled = false;
+}
+
+void fail_at_allocation_index(size_t failure_index) {
+  _allocation_failure_index = failure_index;
+}
+
+set<size_t> outstanding_allocation_indices() {
+  set<size_t> result;
+  for (const auto &entry : _outstanding_allocations) {
+    result.insert(entry.second);
+  }
+  return result;
+}
+
+size_t allocation_count() {
+  return _allocation_count;
+}
+
+}  // namespace record_alloc
+
+static void *record_allocation(void *result) {
+  if (!_enabled)
+    return result;
+
+  if (_allocation_count > _allocation_failure_index) {
+    free(result);
+    Assert::Failure("Allocated after a previous allocation failed!");
+  }
+
+  if (_allocation_count == _allocation_failure_index) {
+    _allocation_count++;
+    free(result);
+    return nullptr;
+  }
+
+  _outstanding_allocations[result] = _allocation_count;
+  _allocation_count++;
+  return result;
+}
+
+static void record_deallocation(void *pointer) {
+  if (!_enabled)
+    return;
+
+  auto entry = _outstanding_allocations.find(pointer);
+  if (entry != _outstanding_allocations.end()) {
+    _outstanding_allocations.erase(entry);
+  }
+}
+
+extern "C" {
+
+void *ts_record_malloc(size_t size) {
+  return record_allocation(malloc(size));
+}
+
+void *ts_record_realloc(void *pointer, size_t size) {
+  record_deallocation(pointer);
+  return record_allocation(realloc(pointer, size));
+}
+
+void *ts_record_calloc(size_t count, size_t size) {
+  return record_allocation(calloc(count, size));
+}
+
+void ts_record_free(void *pointer) {
+  free(pointer);
+  record_deallocation(pointer);
+}
+
+}

spec/helpers/record_alloc.h

@@ -0,0 +1,16 @@
+#ifndef HELPERS_RECORD_ALLOC_H_
+#define HELPERS_RECORD_ALLOC_H_
+
+#include <set>
+
+namespace record_alloc {
+
+void start();
+void stop();
+void fail_at_allocation_index(size_t failure_index);
+std::set<size_t> outstanding_allocation_indices();
+size_t allocation_count();
+
+}  // namespace record_alloc
+
+#endif  // HELPERS_RECORD_ALLOC_H_

spec/helpers/test_languages.cc

@@ -16,9 +16,9 @@ int libcompiler_mtime = -1;
 
 const char libcompiler_path[] =
 #if defined(__linux)
-  "out/Debug/obj.target/libcompiler.a"
+  "out/Test/obj.target/libcompiler.a"
 #else
-  "out/Debug/libcompiler.a"
+  "out/Test/libcompiler.a"
 #endif
 ;
 

spec/runtime/parser_spec.cc

@@ -2,6 +2,7 @@
 #include "helpers/spy_input.h"
 #include "helpers/test_languages.h"
 #include "helpers/log_debugger.h"
+#include "helpers/record_alloc.h"
 
 START_TEST
 
@@ -14,11 +15,14 @@ describe("Parser", [&]() {
   before_each([&]() {
     chunk_size = 3;
     input = nullptr;
+
     doc = ts_document_make();
   });
 
   after_each([&]() {
-    ts_document_free(doc);
+    if (doc)
+      ts_document_free(doc);
+
     if (input)
       delete input;
   });
@@ -421,6 +425,32 @@ describe("Parser", [&]() {
       AssertThat(ts_node_end_byte(root), Equals(strlen("'\u03A9\u03A9\u03A9 \u2014 \u0394\u0394';")));
     });
   });
+
+  describe("handling allocation failures", [&]() {
+    before_each([&]() {
+      record_alloc::start();
+    });
+
+    after_each([&]() {
+      record_alloc::stop();
+    });
+
+    it("handles failures when allocating documents", [&]() {
+      TSDocument *document = ts_document_make();
+      ts_document_free(document);
+      AssertThat(record_alloc::outstanding_allocation_indices(), IsEmpty());
+
+      size_t allocation_count = record_alloc::allocation_count();
+      AssertThat(allocation_count, IsGreaterThan<size_t>(1));
+
+      for (size_t i = 0; i < allocation_count; i++) {
+        record_alloc::start();
+        record_alloc::fail_at_allocation_index(i);
+        AssertThat(ts_document_make(), Equals<TSDocument *>(nullptr));
+        AssertThat(record_alloc::outstanding_allocation_indices(), IsEmpty());
+      }
+    });
+  });
 });
 
 END_TEST

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) {

tests.gyp

@@ -23,8 +23,8 @@
       'libraries': [
         '-ldl'
       ],
-      'default_configuration': 'Debug',
-      'configurations': {'Debug': {}, 'Release': {}},
+      'default_configuration': 'Test',
+      'configurations': {'Test': {}, 'Release': {}},
       'cflags': [
         '-g',
         '-O0',