diff --git a/CMakeLists.txt b/CMakeLists.txt
index 0c43c5c494..e0f05e1205 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -408,6 +408,19 @@ main(void)
 if(TS_HAS_SET_MATCH_LIMIT)
     set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -DNVIM_TS_HAS_SET_MATCH_LIMIT")
 endif()
+check_c_source_compiles("
+#include <stdlib.h>
+#include <tree_sitter/api.h>
+int
+main(void)
+{
+  ts_set_allocator(malloc, calloc, realloc, free);
+  return 0;
+}
+" TS_HAS_SET_ALLOCATOR)
+if(TS_HAS_SET_ALLOCATOR)
+  set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -DNVIM_TS_HAS_SET_ALLOCATOR")
+endif()
 
 # Note: The test lib requires LuaJIT; it will be skipped if LuaJIT is missing.
 option(PREFER_LUA "Prefer Lua over LuaJIT in the nvim executable." OFF)
diff --git a/src/nvim/lua/treesitter.c b/src/nvim/lua/treesitter.c
index 60a000843f..f4067ad02f 100644
--- a/src/nvim/lua/treesitter.c
+++ b/src/nvim/lua/treesitter.c
@@ -129,6 +129,10 @@ void tslua_init(lua_State *L)
   build_meta(L, TS_META_QUERY, query_meta);
   build_meta(L, TS_META_QUERYCURSOR, querycursor_meta);
   build_meta(L, TS_META_TREECURSOR, treecursor_meta);
+
+#ifdef NVIM_TS_HAS_SET_ALLOCATOR
+  ts_set_allocator(xmalloc, xcalloc, xrealloc, xfree);
+#endif
 }
 
 int tslua_has_language(lua_State *L)