BUGFIX: Making sure the array index is locked

pavanky · pavanky · commit 0cee88bcfeff · 2015-03-23T11:45:04.000-04:00
- also added support for logical indexing
diff --git a/src/api/c/assign.cpp b/src/api/c/assign.cpp
@@ -190,8 +190,9 @@ af_err af_assign_gen(af_array *out,
         }
 
         for (int i = 0; i < 4; i++) {
-            if (oDims[i] != rhsDims[i])
+            if (oDims[i] != rhsDims[i]) {
                 AF_ERROR("Size mismatch between input and output", AF_ERR_SIZE);
+            }
         }
 
         af_index_t idxrs[4];
diff --git a/src/api/cpp/array.cpp b/src/api/cpp/array.cpp
@@ -16,10 +16,24 @@
 #include <af/index.h>
 #include <af/device.h>
 #include <af/gfor.h>
+#include <af/algorithm.h>
 #include "error.hpp"
 
 namespace af
 {
+    static void copyIndices(af_index_t inds[4], af_index_t indices[4])
+    {
+        for (int i = 0; i < 4; i++) {
+            if (!indices[i].mIsSeq) {
+                AF_THROW(af_weak_copy(&inds[i].mIndexer.arr, indices[i].mIndexer.arr));
+            } else {
+                inds[i].mIndexer.seq = indices[i].mIndexer.seq;
+            }
+            inds[i].mIsSeq = indices[i].mIsSeq;
+            inds[i].isBatch = indices[i].isBatch;
+        }
+    }
+
     static af_index_t toIndices(const seq &s)
     {
         af_index_t res;
@@ -29,15 +43,31 @@ namespace af
         return res;
     }
 
-    static af_index_t toIndices(const array &idx)
+    static af_index_t toIndices(const array &idx0)
     {
         af_index_t res;
-        res.mIndexer.arr = idx.get();
+
+        array idx = idx0.isbool() ? where(idx0) : idx0;
+        af_array arr = 0;
+        AF_THROW(af_weak_copy(&arr, idx.get()));
+        res.mIndexer.arr = arr;
+
         res.mIsSeq = false;
         res.isBatch = false;
         return res;
     }
 
+    void cleanIndices(af_index_t indices[4])
+    {
+        for (int i = 0; i < 4; i++) {
+            if (!indices[i].mIsSeq) {
+                AF_THROW(af_destroy_array(indices[i].mIndexer.arr));
+            }
+            // Just to be safe
+            indices[i] = toIndices(span);
+        }
+    }
+
     static int gforDim(af_index_t seqs[4])
     {
         for (int i = 0; i < 4; i++) {
@@ -246,6 +276,8 @@ namespace af
             arr = temp;
         }
 
+        cleanIndices(indices);
+
         isRef = false;
         AF_THROW(err);
         return arr;
@@ -497,12 +529,12 @@ namespace af
 
             af_array tmp;
             AF_THROW(af_assign_gen(&tmp, arr, nd, indices, other_arr));
-            parent->set(tmp);
+            cleanIndices(indices);
 
+            parent->set(tmp);
             if (dim >= 0 && is_reordered) AF_THROW(af_destroy_array(other_arr));
 
             isRef = false;
-
         } else {
 
             if (this->get() == other.get()) {
@@ -527,13 +559,17 @@ namespace af
             af_array lhs;                                               \
             int dim = gforDim(this->indices);                           \
             AF_THROW(af_weak_copy(&lhs, this->arr));                    \
+            af_index_t inds[4];                                         \
+            /*FIXME: Figure out a way to not perform the copy*/         \
+            copyIndices(inds, indices);                                 \
             unsigned ndims = numDims(lhs);                              \
             /* FIXME: Unify with other af_assign_gen */                 \
             array tmp = *this op1 other;                                \
             af_array tmp_arr = tmp.get();                               \
             af_array out = 0;                                           \
             tmp_arr = (dim == -1) ? tmp_arr : gforReorder(tmp_arr, dim); \
-            AF_THROW(af_assign_gen(&out, lhs, ndims, indices, tmp_arr)); \
+            AF_THROW(af_assign_gen(&out, lhs, ndims, inds, tmp_arr));   \
+            cleanIndices(indices);                                      \
             AF_THROW(af_destroy_array(this->arr));                      \
             if (dim >= 0) AF_THROW(af_destroy_array(tmp_arr));          \
             this->arr = lhs;                                            \

Original file line number	Diff line number	Diff line change
`@@ -190,8 +190,9 @@ af_err af_assign_gen(af_array *out,`
`190`	`190`	`}`
`191`	`191`
`192`	`192`	`for (int i = 0; i < 4; i++) {`
`193`		`- if (oDims[i] != rhsDims[i])`
	`193`	`+ if (oDims[i] != rhsDims[i]) {`
`194`	`194`	`AF_ERROR("Size mismatch between input and output", AF_ERR_SIZE);`
	`195`	`+ }`
`195`	`196`	`}`
`196`	`197`
`197`	`198`	`af_index_t idxrs[4];`