Fix index copy constructor by retaining array based index

umar456 · umar456 · commit c3aad36aa590 · 2020-11-11T09:48:20.000-05:00
Array based indexing caused segfaults sometimes when the index object was not
immidiately consumed by the indexing operations. This PR fixes this by
retaining the index array on the copy constructor
diff --git a/src/api/cpp/index.cpp b/src/api/cpp/index.cpp
@@ -66,7 +66,13 @@ index::index(const af::array &idx0) : impl{} {
     impl.isBatch = false;
 }
 
-index::index(const af::index &idx0) : impl{idx0.impl} {}  // NOLINT
+index::index(const af::index &idx0) : impl{idx0.impl} {
+    if (!impl.isSeq && impl.idx.arr) {
+        // increment reference count to avoid double free
+        // when/if idx0 is destroyed
+        AF_THROW(af_retain_array(&impl.idx.arr, impl.idx.arr));
+    }
+}
 
 // NOLINTNEXTLINE(hicpp-noexcept-move, performance-noexcept-move-constructor)
 index::index(index &&idx0) : impl{idx0.impl} { idx0.impl.idx.arr = nullptr; }
@@ -79,7 +85,7 @@ index &index::operator=(const index &idx0) {
     if (this == &idx0) { return *this; }
 
     impl = idx0.get();
-    if (!impl.isSeq) {
+    if (!impl.isSeq && impl.idx.arr) {
         // increment reference count to avoid double free
         // when/if idx0 is destroyed
         AF_THROW(af_retain_array(&impl.idx.arr, impl.idx.arr));
diff --git a/test/index.cpp b/test/index.cpp
@@ -1764,6 +1764,19 @@ TEST(Index, ISSUE_2273_Flipped) {
     ASSERT_ARRAYS_EQ(input_slice_gold, input_slice);
 }
 
+TEST(Index, CopiedIndexDestroyed) {
+    array in = randu(10, 10);
+    array a  = constant(1, 10);
+
+    af::index index1(a);
+    af::index index2(seq(10));
+
+    af::index index3(index1);
+    { af::index index4(index1); }
+
+    af_print(in(index1, index2));
+}
+
 // clang-format off
 class IndexDocs : public ::testing::Test {
 public:
