Fix complex matrix assignment ownership

Signed-off-by: Minh Vu <vuhoangminh97@gmail.com>

Author: fallintoplace

runtime/cudaq/operators/matrix.h

@@ -20,6 +20,7 @@
 #include <complex>
 #include <iterator>
 #include <optional>
+#include <utility>
 #include <vector>
 
 namespace Eigen {
@@ -74,14 +75,17 @@ class complex_matrix {
 
   complex_matrix(const complex_matrix &other)
       : dimensions{other.dimensions},
-        data{new value_type[get_size(other.dimensions)]},
+        data{get_size(other.dimensions)
+                 ? new value_type[get_size(other.dimensions)]
+                 : nullptr},
         internal_order(other.internal_order) {
-    std::copy(other.data, other.data + get_size(dimensions), data);
+    if (data)
+      std::copy(other.data, other.data + get_size(dimensions), data);
   }
 
   complex_matrix(const complex_matrix &other, order order);
 
-  complex_matrix(complex_matrix &&other)
+  complex_matrix(complex_matrix &&other) noexcept
       : dimensions{other.dimensions}, data{other.data},
         internal_order(other.internal_order) {
     other.data = nullptr;
@@ -96,18 +100,14 @@ class complex_matrix {
   }
 
   complex_matrix &operator=(const complex_matrix &other) {
-    dimensions = other.dimensions;
-    data = new value_type[get_size(other.dimensions)];
-    std::copy(other.data, other.data + get_size(dimensions), data);
-    internal_order = other.internal_order;
+    complex_matrix copy(other);
+    swap(copy);
     return *this;
   }
 
-  complex_matrix &operator=(complex_matrix &&other) {
-    dimensions = other.dimensions;
-    data = other.data;
-    other.data = nullptr;
-    internal_order = other.internal_order;
+  complex_matrix &operator=(complex_matrix &&other) noexcept {
+    complex_matrix moved(std::move(other));
+    swap(moved);
     return *this;
   }
 
@@ -232,6 +232,13 @@ class complex_matrix {
 
   static void check_size(std::size_t size, const Dimensions &dim);
 
+  void swap(complex_matrix &other) noexcept {
+    using std::swap;
+    swap(dimensions, other.dimensions);
+    swap(data, other.data);
+    swap(internal_order, other.internal_order);
+  }
+
   void swap(complex_matrix::value_type *new_data) {
     if (data)
       delete[] data;

unittests/utils/Matrix.cpp

@@ -55,6 +55,39 @@ TEST(Tensor, initializationError) {
   }
 }
 
+TEST(Tensor, assignment) {
+  const cudaq::complex_matrix empty;
+  const cudaq::complex_matrix small({1., 2., 3., 4.});
+  const cudaq::complex_matrix large({1., 2., 3., 4., 5., 6.}, {2, 3},
+                                    cudaq::complex_matrix::order::column_major);
+
+  cudaq::complex_matrix copy;
+  for (std::size_t i = 0; i < 100; ++i) {
+    copy = empty;
+    EXPECT_EQ(copy, empty);
+    copy = small;
+    EXPECT_EQ(copy, small);
+    copy = large;
+    EXPECT_EQ(copy, large);
+  }
+
+  copy = copy;
+  EXPECT_EQ(copy, large);
+
+  cudaq::complex_matrix moved({7., 8., 9., 10.});
+  for (std::size_t i = 0; i < 100; ++i) {
+    moved = cudaq::complex_matrix();
+    EXPECT_EQ(moved, empty);
+    moved = cudaq::complex_matrix(small);
+    EXPECT_EQ(moved, small);
+    moved = cudaq::complex_matrix(large);
+    EXPECT_EQ(moved, large);
+  }
+
+  moved = std::move(moved);
+  EXPECT_EQ(moved, large);
+}
+
 TEST(Tensor, access) {
   {
     cudaq::complex_matrix m1({1., 0., 0., 1.});