From e742186c25f55349313ce6c50ae91a8744b5ee56 Mon Sep 17 00:00:00 2001
From: Mingkai Dong <mingkaidong@gmail.com>
Date: Fri, 17 Dec 2021 19:53:02 +0800
Subject: [PATCH] Allow to change the allocator dynamically

---
 lib/src/alloc.c | 72 +++++++++++++++++++++++++++++++++++++++++++++++++
 lib/src/alloc.h | 70 ++++++++---------------------------------------
 lib/src/lib.c   |  1 +
 3 files changed, 84 insertions(+), 59 deletions(-)
 create mode 100644 lib/src/alloc.c

diff --git a/lib/src/alloc.c b/lib/src/alloc.c
new file mode 100644
index 00000000..5ffc050c
--- /dev/null
+++ b/lib/src/alloc.c
@@ -0,0 +1,72 @@
+#include "alloc.h"
+
+#if defined(TREE_SITTER_ALLOCATION_TRACKING)
+
+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 *);
+bool ts_toggle_allocation_recording(bool);
+
+static struct allocator 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;
+
+#else
+
+#include <stdlib.h>
+
+static inline bool ts_toggle_allocation_recording(bool value) {
+  (void)value;
+  return false;
+}
+
+
+static inline void *ts_malloc_default(size_t size) {
+  void *result = malloc(size);
+  if (size > 0 && !result) {
+    fprintf(stderr, "tree-sitter failed to allocate %zu bytes", size);
+    exit(1);
+  }
+  return result;
+}
+
+static inline void *ts_calloc_default(size_t count, size_t size) {
+  void *result = calloc(count, size);
+  if (count > 0 && !result) {
+    fprintf(stderr, "tree-sitter failed to allocate %zu bytes", count * size);
+    exit(1);
+  }
+  return result;
+}
+
+static inline void *ts_realloc_default(void *buffer, size_t size) {
+  void *result = realloc(buffer, size);
+  if (size > 0 && !result) {
+    fprintf(stderr, "tree-sitter failed to reallocate %zu bytes", size);
+    exit(1);
+  }
+  return result;
+}
+
+static inline void ts_free_default(void *buffer) {
+  free(buffer);
+}
+
+static struct allocator 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
+struct allocator *ts_allocator = &ts_default_allocator;
+
+
+#endif
diff --git a/lib/src/alloc.h b/lib/src/alloc.h
index dd487ca2..131a8b49 100644
--- a/lib/src/alloc.h
+++ b/lib/src/alloc.h
@@ -9,75 +9,27 @@ extern "C" {
 #include <stdbool.h>
 #include <stdio.h>
 
-#if defined(TREE_SITTER_ALLOCATION_TRACKING)
+struct allocator {
+  void *(*malloc)(size_t);
+  void *(*calloc)(size_t, size_t);
+  void *(*realloc)(void *, size_t);
+  void (*free)(void *);
+};
 
-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 *);
-bool ts_toggle_allocation_recording(bool);
-
-#define ts_malloc  ts_record_malloc
-#define ts_calloc  ts_record_calloc
-#define ts_realloc ts_record_realloc
-#define ts_free    ts_record_free
-
-#else
+extern struct allocator *ts_allocator;
 
 // Allow clients to override allocation functions
-
 #ifndef ts_malloc
-#define ts_malloc  ts_malloc_default
+#define ts_malloc  ts_allocator->malloc
 #endif
 #ifndef ts_calloc
-#define ts_calloc  ts_calloc_default
+#define ts_calloc  ts_allocator->calloc
 #endif
 #ifndef ts_realloc
-#define ts_realloc ts_realloc_default
+#define ts_realloc ts_allocator->realloc
 #endif
 #ifndef ts_free
-#define ts_free    ts_free_default
-#endif
-
-#include <stdlib.h>
-
-static inline bool ts_toggle_allocation_recording(bool value) {
-  (void)value;
-  return false;
-}
-
-
-static inline void *ts_malloc_default(size_t size) {
-  void *result = malloc(size);
-  if (size > 0 && !result) {
-    fprintf(stderr, "tree-sitter failed to allocate %zu bytes", size);
-    exit(1);
-  }
-  return result;
-}
-
-static inline void *ts_calloc_default(size_t count, size_t size) {
-  void *result = calloc(count, size);
-  if (count > 0 && !result) {
-    fprintf(stderr, "tree-sitter failed to allocate %zu bytes", count * size);
-    exit(1);
-  }
-  return result;
-}
-
-static inline void *ts_realloc_default(void *buffer, size_t size) {
-  void *result = realloc(buffer, size);
-  if (size > 0 && !result) {
-    fprintf(stderr, "tree-sitter failed to reallocate %zu bytes", size);
-    exit(1);
-  }
-  return result;
-}
-
-static inline void ts_free_default(void *buffer) {
-  free(buffer);
-}
-
+#define ts_free    ts_allocator->free
 #endif
 
 #ifdef __cplusplus
diff --git a/lib/src/lib.c b/lib/src/lib.c
index 289d32f4..5aab20d5 100644
--- a/lib/src/lib.c
+++ b/lib/src/lib.c
@@ -5,6 +5,7 @@
 
 #define _POSIX_C_SOURCE 200112L
 
+#include "./alloc.c"
 #include "./get_changed_ranges.c"
 #include "./language.c"
 #include "./lexer.c"