From db9f6602ca7a292987025dab6f8d28b65ccc99cc Mon Sep 17 00:00:00 2001
From: wrtobin <tobin6@llnl.gov>
Date: Wed, 30 Oct 2019 15:43:32 -0700
Subject: [PATCH] adding test case to catch realloc bug with default
 constructor and test case to ensure realloc on preallocated managed arrays
 copies over data correctly, also implement a fix for the realloc bug

---
 src/chai/ArrayManager.inl               | 18 ++++++++++++------
 src/chai/ManagedArray.inl               |  3 ++-
 tests/unit/managed_array_unit_tests.cpp | 19 +++++++++++++++++++
 3 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/src/chai/ArrayManager.inl b/src/chai/ArrayManager.inl
index 90d4f280..f55954e5 100644
--- a/src/chai/ArrayManager.inl
+++ b/src/chai/ArrayManager.inl
@@ -47,20 +47,26 @@ void* ArrayManager::reallocate(void* pointer, size_t elems, PointerRecord* point
   for (int space = CPU; space < NUM_EXECUTION_SPACES; ++space) {
     void* old_ptr = pointer_record->m_pointers[space];
 
-    if (old_ptr) {
-      pointer_record->m_user_callback(ACTION_ALLOC, ExecutionSpace(space), sizeof(T) * elems);
-      void* new_ptr = m_allocators[space]->allocate(sizeof(T)*elems);
+    pointer_record->m_user_callback(ACTION_ALLOC, ExecutionSpace(space), sizeof(T) * elems);
+    void* new_ptr = m_allocators[space]->allocate(sizeof(T)*elems);
 
+    if (old_ptr) {
       m_resource_manager.copy(new_ptr, old_ptr, num_bytes_to_copy);
+    }
+
+    pointer_record->m_user_callback(ACTION_FREE, ExecutionSpace(space), sizeof(T) * elems);
 
-      pointer_record->m_user_callback(ACTION_FREE, ExecutionSpace(space), sizeof(T) * elems);
+    if (old_ptr) {
       m_allocators[space]->deallocate(old_ptr);
+    }
 
-      pointer_record->m_pointers[space] = new_ptr;
+    pointer_record->m_pointers[space] = new_ptr;
 
+    if (old_ptr) {
       m_pointer_map.erase(old_ptr);
-      m_pointer_map.insert(new_ptr, pointer_record);
     }
+
+    m_pointer_map.insert(new_ptr, pointer_record);
   }
 
   pointer_record->m_size = sizeof(T) * elems;
diff --git a/src/chai/ManagedArray.inl b/src/chai/ManagedArray.inl
index 87cb6d36..5fd06231 100644
--- a/src/chai/ManagedArray.inl
+++ b/src/chai/ManagedArray.inl
@@ -31,8 +31,9 @@ CHAI_HOST_DEVICE ManagedArray<T>::ManagedArray():
   m_pointer_record->m_user_callback = [](Action, ExecutionSpace, size_t) {};
 
   for (int space = CPU;  space < NUM_EXECUTION_SPACES; space++) {
-    m_pointer_record->m_allocators[space] = 
+    m_pointer_record->m_allocators[space] =
       m_resource_manager->getAllocatorId(ExecutionSpace(space));
+    m_pointer_record->m_owned[space] = true;
   }
 #endif
 }
diff --git a/tests/unit/managed_array_unit_tests.cpp b/tests/unit/managed_array_unit_tests.cpp
index 50fb1696..cdec57a6 100644
--- a/tests/unit/managed_array_unit_tests.cpp
+++ b/tests/unit/managed_array_unit_tests.cpp
@@ -35,6 +35,25 @@ TEST(ManagedArray, SpaceConstructorCPU)
   array.free();
 }
 
+TEST(ManagedArray, ReallocDefaultConstructor)
+{
+  chai::ManagedArray<float> array;
+  array.reallocate(1);
+  float val = array[0];
+  (void)val;
+  ASSERT_EQ(array.size(), 1u);
+}
+
+TEST(ManagedArray, ReallocSizeConstructor)
+{
+  chai::ManagedArray<float> array(10);
+  array[0] = 1.0;
+  array[9] = 2.0;
+  array.reallocate(20);
+  ASSERT_EQ(array[0], 1.0);
+  ASSERT_EQ(array[9], 2.0);
+}
+
 #if defined(CHAI_ENABLE_CUDA) || defined(CHAI_ENABLE_HIP)
 TEST(ManagedArray, SpaceConstructorGPU)
 {