Merge pull request #209 from dominikgrewe/allocator

Pass optional allocator to storage.

1 files changed, 30 insertions, 14 deletions

diff --git a/generic/Storage.c b/generic/Storage.c
index 7de1828..aa8242a 100644
--- a/generic/Storage.c
+++ b/generic/Storage.c
@@ -4,22 +4,32 @@
 
 static int torch_Storage_(new)(lua_State *L)
 {
+  int index = 1;
   THStorage *storage;
-  if(lua_type(L, 1) == LUA_TSTRING)
+  THAllocator *allocator = luaT_toudata(L, index, "torch.Allocator");
+  if (allocator) index++;
+
+  if(lua_type(L, index) == LUA_TSTRING)
   {
-    const char *fileName = luaL_checkstring(L, 1);
-    int isShared = luaT_optboolean(L, 2, 0);
-    long size = luaL_optlong(L, 3, 0);
+    if (allocator)
+      THError("Passing allocator not supported when using file mapping");
+
+    const char *fileName = luaL_checkstring(L, index);
+    int isShared = luaT_optboolean(L, index + 1, 0);
+    long size = luaL_optlong(L, index + 2, 0);
     storage = THStorage_(newWithMapping)(fileName, size, isShared);
   }
-  else if(lua_type(L, 1) == LUA_TTABLE)
+  else if(lua_type(L, index) == LUA_TTABLE)
   {
-    long size = lua_objlen(L, 1);
+    long size = lua_objlen(L, index);
     long i;
-    storage = THStorage_(newWithSize)(size);
+    if (allocator)
+      storage = THStorage_(newWithAllocator)(size, allocator, NULL);
+    else
+      storage = THStorage_(newWithSize)(size);
     for(i = 1; i <= size; i++)
     {
-      lua_rawgeti(L, 1, i);
+      lua_rawgeti(L, index, i);
       if(!lua_isnumber(L, -1))
       {
         THStorage_(free)(storage);
@@ -29,17 +39,23 @@ static int torch_Storage_(new)(lua_State *L)
       lua_pop(L, 1);
     }
   }
-  else if(lua_type(L, 2) == LUA_TNUMBER)
+  else if(lua_type(L, index + 1) == LUA_TNUMBER)
   {
-    long size = luaL_optlong(L, 1, 0);
-    real *ptr = (real *)luaL_optlong(L, 2, 0);
-    storage = THStorage_(newWithData)(ptr, size);
+    long size = luaL_optlong(L, index, 0);
+    real *ptr = (real *)luaL_optlong(L, index + 1, 0);
+    if (allocator)
+      storage = THStorage_(newWithDataAndAllocator)(ptr, size, allocator, NULL);
+    else
+      storage = THStorage_(newWithData)(ptr, size);
     storage->flag = 0;
   }
   else
   {
-    long size = luaL_optlong(L, 1, 0);
-    storage = THStorage_(newWithSize)(size);
+    long size = luaL_optlong(L, index, 0);
+    if (allocator)
+      storage = THStorage_(newWithAllocator)(size, allocator, NULL);
+    else
+      storage = THStorage_(newWithSize)(size);
   }
   luaT_pushudata(L, storage, torch_Storage);
   return 1;