diff --git a/lib/include/tree_sitter/api.h b/lib/include/tree_sitter/api.h
index 66ceaea1..f1192e34 100644
--- a/lib/include/tree_sitter/api.h
+++ b/lib/include/tree_sitter/api.h
@@ -134,10 +134,31 @@ typedef enum {
   TSQueryErrorLanguage,
 } TSQueryError;
 
+typedef struct {
+  void *(*malloc)(size_t);
+  void *(*calloc)(size_t, size_t);
+  void *(*realloc)(void *, size_t);
+  void (*free)(void *);
+} TSAllocator;
+
 /********************/
 /* Section - Parser */
 /********************/
 
+
+/**
+ * The allocator.
+ */
+extern TSAllocator *ts_allocator;
+
+/**
+ * Switch to a new allocator.
+ *
+ * Returns the old allocator. The caller needs to take care of any allocated
+ * memory managed by the old allocator.
+ */
+TSAllocator *ts_set_allocator(TSAllocator *new_alloc);
+
 /**
  * Create a new parser.
  */
diff --git a/lib/src/alloc.c b/lib/src/alloc.c
index 6f62a777..391430f7 100644
--- a/lib/src/alloc.c
+++ b/lib/src/alloc.c
@@ -8,14 +8,14 @@ void *ts_record_realloc(void *, size_t);
 void ts_record_free(void *);
 bool ts_toggle_allocation_recording(bool);
 
-static struct allocator ts_tracking_allocator = {
+static TSAllocator ts_tracking_allocator = {
   .malloc = ts_record_malloc,
   .calloc = ts_record_calloc,
   .realloc = ts_record_realloc,
   .free = ts_record_free,
 };
 
-struct allocator *ts_allocator = &ts_tracking_allocator;
+TSAllocator *ts_allocator = &ts_tracking_allocator;
 
 #else
 
@@ -58,7 +58,7 @@ static inline void ts_free_default(void *buffer) {
   free(buffer);
 }
 
-static struct allocator ts_default_allocator = {
+static TSAllocator ts_default_allocator = {
   .malloc = ts_malloc_default,
   .calloc = ts_calloc_default,
   .realloc = ts_realloc_default,
@@ -66,7 +66,14 @@ static struct allocator ts_default_allocator = {
 };
 
 // Allow clients to override allocation functions dynamically
-struct allocator *ts_allocator = &ts_default_allocator;
+TSAllocator *ts_allocator = &ts_default_allocator;
 
+#endif // defined(TREE_SITTER_ALLOCATION_TRACKING)
+
+TSAllocator *ts_set_allocator(TSAllocator *new_alloc)
+{
+	TSAllocator *old = ts_allocator;
+	ts_allocator = new_alloc;
+	return old;
+}
 
-#endif
diff --git a/lib/src/alloc.h b/lib/src/alloc.h
index 35ba77d2..6c04ae18 100644
--- a/lib/src/alloc.h
+++ b/lib/src/alloc.h
@@ -1,6 +1,8 @@
 #ifndef TREE_SITTER_ALLOC_H_
 #define TREE_SITTER_ALLOC_H_
 
+#include "tree_sitter/api.h"
+
 #ifdef __cplusplus
 extern "C" {
 #endif
@@ -9,14 +11,7 @@ extern "C" {
 #include <stdbool.h>
 #include <stdio.h>
 
-struct allocator {
-  void *(*malloc)(size_t);
-  void *(*calloc)(size_t, size_t);
-  void *(*realloc)(void *, size_t);
-  void (*free)(void *);
-};
-
-extern struct allocator *ts_allocator;
+extern TSAllocator *ts_allocator;
 
 // Allow clients to override allocation functions
 #ifndef ts_malloc