THAlloc/THRealloc attempt lua GC on failure

Author: adamlerer

THGeneral.c

@@ -13,6 +13,9 @@ static void defaultTorchErrorHandlerFunction(const char *msg, void *data)
 static __thread void (*torchErrorHandlerFunction)(const char *msg, void *data) = defaultTorchErrorHandlerFunction;
 static __thread void *torchErrorHandlerData;
 
+static __thread void (*torchGCFunction)(void *data) = NULL;
+static __thread void *torchGCData;
+
 void _THError(const char *file, const int line, const char *fmt, ...)
 {
  char msg[2048];
@@ -91,15 +94,15 @@ void THSetArgErrorHandler( void (*torchArgErrorHandlerFunction_)(int argNumber,
  torchArgErrorHandlerData = data;
 }
 
-void* THAlloc(long size)
+void THSetGCHandler( void (*torchGCFunction_)(void *data), void *data )
 {
- void *ptr;
-
- if(size < 0)
- THError("$ Torch: invalid memory size -- maybe an overflow?");
+ torchGCFunction = torchGCFunction_;
+ torchGCData = data;
+}
 
- if(size == 0)
- return NULL;
+void* THAllocInternal(long size)
+{
+ void *ptr;
 
  if (size > 5120)
  {
@@ -116,6 +119,25 @@ void* THAlloc(long size)
  {
  ptr = malloc(size);
  }
+ return ptr;
+}
+
+void* THAlloc(long size)
+{
+ void *ptr;
+
+ if(size < 0)
+ THError("$ Torch: invalid memory size -- maybe an overflow?");
+
+ if(size == 0)
+ return NULL;
+
+ ptr = THAllocInternal(size);
+
+ if(!ptr && torchGCFunction) {
+ torchGCFunction(torchGCData);
+ ptr = THAllocInternal(size);
+ }
 
  if(!ptr)
  THError("$ Torch: not enough memory: you tried to allocate %dGB. Buy new RAM!", size/1073741824);
@@ -137,10 +159,17 @@ void* THRealloc(void *ptr, long size)
  if(size < 0)
  THError("$ Torch: invalid memory size -- maybe an overflow?");
 
- ptr = realloc(ptr, size);
- if(!ptr)
+ void *newptr = realloc(ptr, size);
+
+ if(!newptr && torchGCFunction) {
+ torchGCFunction(torchGCData);
+ newptr = realloc(ptr, size);
+ }
+
+ if(!newptr)
  THError("$ Torch: not enough memory: you tried to reallocate %dGB. Buy new RAM!", size/1073741824);
- return ptr;
+
+ return newptr;
 }
 
 void THFree(void *ptr)

THGeneral.h.in

@@ -52,6 +52,7 @@ TH_API void THSetArgErrorHandler( void (*torchArgErrorHandlerFunction)(int argNu
 TH_API void* THAlloc(long size);
 TH_API void* THRealloc(void *ptr, long size);
 TH_API void THFree(void *ptr);
+TH_API void THSetGCHandler( void (*torchGCHandlerFunction)(void *data), void *data );
 
 #define THError(...) _THError(__FILE__, __LINE__, __VA_ARGS__)
 #define THArgCheck(...) _THArgCheck(__FILE__, __LINE__, __VA_ARGS__)