Adding support for Storage views to cutorch.

1 files changed, 19 insertions, 2 deletions

diff --git a/torch/generic/Storage.c b/torch/generic/Storage.c
index 995e57f..9dc1333 100644
--- a/torch/generic/Storage.c
+++ b/torch/generic/Storage.c
@@ -30,12 +30,29 @@ static int torch_Storage_(new)(lua_State *L)
       lua_pop(L, 1);
     }
   }
+  else if(lua_type(L, 1) == LUA_TUSERDATA)
+  {
+    THStorage *src = luaT_checkudata(L, 1, torch_Storage);
+    real *ptr = src->data;
+    long offset = luaL_optlong(L, 2, 1) - 1;
+    if (offset < 0 || offset >= src->size) {
+      luaL_error(L, "offset out of bounds");
+    }
+    long size = luaL_optlong(L, 3, src->size - offset);
+    if (size < 1 || size > (src->size - offset)) {
+      luaL_error(L, "size out of bounds");
+    }
+    storage = THStorage_(newWithData)(state, ptr + offset, size);
+    storage->flag = TH_STORAGE_REFCOUNTED | TH_STORAGE_VIEW;
+    storage->view = src;
+    THStorage_(retain)(storage->view);
+  }
   else if(lua_type(L, 2) == LUA_TNUMBER)
   {
     long size = luaL_optlong(L, 1, 0);
     real *ptr = (real *)luaL_optlong(L, 2, 0);
     storage = THStorage_(newWithData)(state, ptr, size);
-    storage->flag = 0;
+    storage->flag = TH_STORAGE_REFCOUNTED;
   }
   else
   {
@@ -183,7 +200,7 @@ static int torch_Storage_(write)(lua_State *L)
 {
   THStorage *storage = luaT_checkudata(L, 1, torch_Storage);
   THFile *file = luaT_checkudata(L, 2, "torch.File");
- 
+
   THFile_writeLongScalar(file, storage->size);
   THFile_writeRealRaw(file, storage->data, storage->size);