improved serialization such that it does not mess up with luajit's GC

2 files changed, 40 insertions, 21 deletions

diff --git a/serialize.lua b/serialize.lua
index 1f24645..39b14b7 100644
--- a/serialize.lua
+++ b/serialize.lua
@@ -1,29 +1,38 @@
 local ffi = require 'ffi'
 local C = ffi.C
 
+require 'torch'
+
 ffi.cdef[[
 void free(void *ptr);
 void *malloc(size_t size);
+THCharStorage* THCharStorage_newWithData(const char *data, long size);
+void THCharStorage_clearFlag(THCharStorage *storage, const char flag);
 ]]
 
-require 'torch'
-
 local serialize = {}
 
 function serialize.save(func)
-   local code = torch.serialize(func) -- DEBUG: make it work without torch too ;)
-   local sz = #code
-   local code_p = ffi.cast('char*', C.malloc(sz)) -- C.malloc(sz+1))
-   assert(code_p ~= nil, 'allocation error during serialization')
---   code_p[sz] = 0
-   ffi.copy(code_p, ffi.cast('const char*', code), sz)
+   local f = torch.MemoryFile()
+   f:binary()
+   f:writeObject(func)
+   local storage = f:storage()
+   local code_p = storage:data()
+   local sz = storage:size()
+    -- refcounted, but do not free mem
+   C.THCharStorage_clearFlag(storage:cdata(), 4)
+   f:close()
    return code_p, sz
 end
 
 function serialize.load(code_p, sz)
-   local code = ffi.string(code_p, sz)
-   C.free(ffi.cast('void*', code_p))
-   return torch.deserialize(code)
+   local storage_p = C.THCharStorage_newWithData(code_p, sz)
+   local storage = torch.pushudata(storage_p, 'torch.CharStorage')
+   local f = torch.MemoryFile(storage)
+   f:binary()
+   local func = f:readObject()
+   f:close()
+   return func
 end
 
 return serialize
diff --git a/sharedserialize.lua b/sharedserialize.lua
index 4f18120..b151cec 100644
--- a/sharedserialize.lua
+++ b/sharedserialize.lua
@@ -1,13 +1,15 @@
 local ffi = require 'ffi'
 local C = ffi.C
 
+require 'torch'
+
 ffi.cdef[[
 void free(void *ptr);
 void *malloc(size_t size);
+THCharStorage* THCharStorage_newWithData(const char *data, long size);
+void THCharStorage_clearFlag(THCharStorage *storage, const char flag);
 ]]
 
-require 'torch'
-
 local serialize = {}
 
 local tensor = {}
@@ -74,11 +76,15 @@ function serialize.save(func)
    sharewrite()
    local status, code_p, sz = pcall(
       function()
-         local code = torch.serialize(func) -- DEBUG: make it work without torch too ;)
-         local sz = #code
-         local code_p = ffi.cast('char*', C.malloc(sz)) -- C.malloc(sz+1))
-         assert(code_p ~= nil, 'allocation error during serialization')
-         ffi.copy(code_p, ffi.cast('const char*', code), sz)
+         local f = torch.MemoryFile()
+         f:binary()
+         f:writeObject(func)
+         local storage = f:storage()
+         local code_p = storage:data()
+         local sz = storage:size()
+         -- refcounted, but do not free mem
+         C.THCharStorage_clearFlag(storage:cdata(), 4)
+         f:close()
          return code_p, sz
       end
    )
@@ -93,9 +99,13 @@ function serialize.load(code_p, sz)
    shareread()
    local status, func = pcall(
       function()
-         local code = ffi.string(code_p, sz)
-         C.free(ffi.cast('void*', code_p))
-         return torch.deserialize(code)
+         local storage_p = C.THCharStorage_newWithData(code_p, sz)
+         local storage = torch.pushudata(storage_p, 'torch.CharStorage')
+         local f = torch.MemoryFile(storage)
+         f:binary()
+         local func = f:readObject()
+         f:close()
+         return func
       end
    )
    unshareread()