add matrix indexing support

1 files changed, 85 insertions, 8 deletions

diff --git a/matrix.c b/matrix.c
index 9d93dba..a58377f 100644
--- a/matrix.c
+++ b/matrix.c
@@ -1,32 +1,44 @@
 #include "common.h"
 
 typedef struct Matrix {
-    long stride;       /* size of a row */
+    long stride;              /* size of a row */
     long ncol, nrow, nmax;    /* dimension of the matrix */
     union {
         float *f;
         double *d;
-    } storage;          /* pointer to actual storage */
+    } data;                   /* pointer to actual storage */
+    long *data_ref;
 } Matrix;
 
 const char *float_matrix_tname = "nerv.FloatMatrix";
 const char *matrix_tname = "nerv.Matrix";
 
+void float_matrix_data_free(Matrix *self) {
+    if (--(*self->data_ref) == 0)
+        free(self->data.f);
+}
+
+void float_matrix_data_retain(Matrix *self) {
+    (*self->data_ref)++;
+}
+
 int float_matrix_new(lua_State *L) {
     Matrix *self = (Matrix *)malloc(sizeof(Matrix));
     self->nrow = luaL_checkinteger(L, 1);
     self->ncol = luaL_checkinteger(L, 2);
     self->nmax = self->nrow * self->ncol;
-    self->stride = sizeof(float) * self->nrow;
-    self->storage.f = (float *)malloc(self->stride * self->ncol);
+    self->stride = sizeof(float) * self->ncol;
+    self->data.f = (float *)malloc(self->stride * self->nrow);
+    self->data_ref = (long *)malloc(sizeof(long));
+    *self->data_ref = 0;
+    float_matrix_data_retain(self);
     luaT_pushudata(L, self, float_matrix_tname);
     return 1;
 }
 
 int float_matrix_destroy(lua_State *L) {
     Matrix *self = luaT_checkudata(L, 1, float_matrix_tname);
-    free(self->storage.f);
-    fprintf(stderr, "[debug] destroyted\n");
+    float_matrix_data_free(self);
     return 0;
 }
 
@@ -35,7 +47,7 @@ int nerv_float_matrix_get_elem(lua_State *L) {
     int idx = luaL_checkinteger(L, 2);
     if (idx < 0 || idx >= self->nmax)
         nerv_error(L, "index must be within range [0, %d)", self->nmax);
-    lua_pushnumber(L, self->storage.f[idx]);
+    lua_pushnumber(L, self->data.f[idx]);
     return 1;
 }
 
@@ -46,10 +58,73 @@ int nerv_float_matrix_set_elem(lua_State *L) {
     long upper = self->nrow * self->ncol;
     if (idx < 0 || idx >= self->nmax)
         nerv_error(L, "index must be within range [0, %d)", self->nmax);
-    self->storage.f[idx] = v;
+    self->data.f[idx] = v;
     return 0;
 }
 
+static Matrix *nerv_float_matrix_getrow(Matrix *self, int row) {
+    Matrix *prow = (Matrix *)malloc(sizeof(Matrix));
+    prow->ncol = self->ncol;
+    prow->nrow = 1;
+    prow->stride = self->stride;
+    prow->nmax = prow->ncol;
+    prow->data.f = (float *)((char *)self->data.f + row * self->stride);
+    prow->data_ref = self->data_ref;
+    float_matrix_data_retain(self);
+    return prow;
+}
+
+static int nerv_float_matrix_newindex(lua_State *L) {
+    Matrix *self = luaT_checkudata(L, 1, float_matrix_tname);
+    if (lua_isnumber(L, 2))
+    {
+        int idx = luaL_checkinteger(L, 2);
+        if (self->nrow == 1)
+        {
+            if (idx < 0 || idx >= self->ncol)
+                nerv_error(L, "index must be within range [0, %d)", self->ncol);
+            self->data.f[idx] = luaL_checknumber(L, 3);
+        }
+        else
+            nerv_error(L, "cannot assign a scalar to row vector");
+        lua_pushboolean(L, 1);
+        return 2;
+    }
+    else
+    {
+        lua_pushboolean(L, 0);
+        return 1;
+    }
+}
+
+
+static int nerv_float_matrix_index(lua_State *L) {
+    Matrix *self = luaT_checkudata(L, 1, float_matrix_tname);
+    if (lua_isnumber(L, 2))
+    {
+        int idx = luaL_checkinteger(L, 2);
+        if (self->nrow == 1)
+        {
+            if (idx < 0 || idx >= self->ncol)
+                nerv_error(L, "index must be within range [0, %d)", self->ncol);
+            lua_pushnumber(L, self->data.f[idx]);
+        }
+        else
+        {
+            if (idx < 0 || idx >= self->nrow)
+                nerv_error(L, "index must be within range [0, %d)", self->nrow);
+            luaT_pushudata(L, nerv_float_matrix_getrow(self, idx), float_matrix_tname);
+        }
+        lua_pushboolean(L, 1);
+        return 2;
+    }
+    else
+    {
+        lua_pushboolean(L, 0);
+        return 1;
+    }
+}
+
 static int nerv_float_matrix_ncol(lua_State *L) {
     Matrix *self = luaT_checkudata(L, 1, float_matrix_tname);
     lua_pushinteger(L, self->ncol);
@@ -68,6 +143,8 @@ static const luaL_Reg float_matrix_methods[] = {
     {"set_elem", nerv_float_matrix_set_elem},
     {"ncol", nerv_float_matrix_ncol},
     {"nrow", nerv_float_matrix_nrow},
+    {"__index__", nerv_float_matrix_index},
+    {"__newindex__", nerv_float_matrix_newindex},
     {NULL, NULL}
 };