diff --git a/lib/include/tree_sitter/api.h b/lib/include/tree_sitter/api.h
index 0b64e60e..4d558a5c 100644
--- a/lib/include/tree_sitter/api.h
+++ b/lib/include/tree_sitter/api.h
@@ -134,13 +134,6 @@ 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 */
 /********************/
@@ -148,15 +141,18 @@ typedef struct {
 /**
  * Switch to a new allocator.
  *
- * Returns true iff the switch successes.
- *
- * This function can only be invoked *once*, before tree-sitter has allocated
- * any memory via malloc/calloc/realloc. Otherwise, memory bugs could occur
- * since some memory is allocated by the old allocator, but freed by the new
- * one.
+ * This function can be invoked more than once. However, the application needs
+ * to pay attention to the memory allocated by the old allocator but might be
+ * freed by the new one.
  *
+ * Specifically, the application either,
+ *  1. ensures all parsers and trees are freed before calling it;
+ *  2. provides an allocator that shares its state with the old allocator.
  */
-bool ts_set_allocator(TSAllocator *new_alloc);
+void ts_set_allocator(void *(*new_malloc)(size_t),
+		      void *(*new_calloc)(size_t, size_t),
+		      void *(*new_realloc)(void *, size_t),
+		      void (*new_free)(void *));
 
 /**
  * Create a new parser.
diff --git a/lib/src/alloc.c b/lib/src/alloc.c
index 65c72f11..9c0458ff 100644
--- a/lib/src/alloc.c
+++ b/lib/src/alloc.c
@@ -8,14 +8,10 @@ void *ts_record_realloc(void *, size_t);
 void ts_record_free(void *);
 bool ts_toggle_allocation_recording(bool);
 
-static TSAllocator ts_tracking_allocator = {
-  .malloc = ts_record_malloc,
-  .calloc = ts_record_calloc,
-  .realloc = ts_record_realloc,
-  .free = ts_record_free,
-};
-
-TSAllocator *ts_allocator = &ts_tracking_allocator;
+void *(*ts_current_malloc)(size_t) = ts_record_malloc;
+void *(*ts_current_calloc)(size_t, size_t) = ts_record_calloc;
+void *(*ts_current_realloc)(void *, size_t) = ts_record_realloc;
+void (*ts_current_free)(void *) = ts_record_free;
 
 #else
 
@@ -58,27 +54,23 @@ static inline void ts_free_default(void *buffer) {
   free(buffer);
 }
 
-static TSAllocator ts_default_allocator = {
-  .malloc = ts_malloc_default,
-  .calloc = ts_calloc_default,
-  .realloc = ts_realloc_default,
-  .free = ts_free_default,
-};
 
 // Allow clients to override allocation functions dynamically
-TSAllocator *ts_allocator = &ts_default_allocator;
+void *(*ts_current_malloc)(size_t) = ts_malloc_default;
+void *(*ts_current_calloc)(size_t, size_t) = ts_calloc_default;
+void *(*ts_current_realloc)(void *, size_t) = ts_realloc_default;
+void (*ts_current_free)(void *) = ts_free_default;
 
 #endif // defined(TREE_SITTER_ALLOCATION_TRACKING)
 
-bool ts_set_allocator(TSAllocator *new_alloc)
+void ts_set_allocator(void *(*new_malloc)(size_t),
+                      void *(*new_calloc)(size_t, size_t),
+                      void *(*new_realloc)(void *, size_t),
+                      void (*new_free)(void *))
 {
-  static bool using_default_allocator = true;
-  if (!using_default_allocator) {
-    fprintf(stderr, "tree-sitter's allocator can only be set once!");
-    return false;
-  }
-  using_default_allocator = false;
-  ts_allocator = new_alloc;
-  return true;
+  ts_current_malloc = new_malloc;
+  ts_current_calloc = new_calloc;
+  ts_current_realloc = new_realloc;
+  ts_current_free = new_free;
 }
 
diff --git a/lib/src/alloc.h b/lib/src/alloc.h
index 6c04ae18..faf7c4ef 100644
--- a/lib/src/alloc.h
+++ b/lib/src/alloc.h
@@ -11,20 +11,23 @@ extern "C" {
 #include <stdbool.h>
 #include <stdio.h>
 
-extern TSAllocator *ts_allocator;
+extern void *(*ts_current_malloc)(size_t);
+extern void *(*ts_current_calloc)(size_t, size_t);
+extern void *(*ts_current_realloc)(void *, size_t);
+extern void (*ts_current_free)(void *);
 
 // Allow clients to override allocation functions
 #ifndef ts_malloc
-#define ts_malloc  ts_allocator->malloc
+#define ts_malloc  ts_current_malloc
 #endif
 #ifndef ts_calloc
-#define ts_calloc  ts_allocator->calloc
+#define ts_calloc  ts_current_calloc
 #endif
 #ifndef ts_realloc
-#define ts_realloc ts_allocator->realloc
+#define ts_realloc ts_current_realloc
 #endif
 #ifndef ts_free
-#define ts_free    ts_allocator->free
+#define ts_free    ts_current_free
 #endif
 
 bool ts_toggle_allocation_recording(bool);