compound assignment operator expressions (+=, -=, *=) work with null …

…DistArrays
ValeevGroup · Sep 30, 2024 · d04fb08 · d04fb08
1 parent ec51edb
commit d04fb08
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 3 deletions.
diff --git a/src/TiledArray/expressions/tsr_expr.h b/src/TiledArray/expressions/tsr_expr.h
@@ -147,7 +147,11 @@ class TsrExpr : public Expr<TsrExpr<Array, Alias>> {
         TiledArray::expressions::is_aliased<D>::value,
         "no_alias() expressions are not allowed on the right-hand side of "
         "the assignment operator.");
-    return operator=(AddExpr<TsrExpr_, D>(*this, other.derived()));
+    if (this->array().is_initialized()) {
+      return operator=(AddExpr<TsrExpr_, D>(*this, other.derived()));
+    } else {
+      return operator=(other);
+    }
   }
 
   /// Expression minus-assignment operator
@@ -160,7 +164,11 @@ class TsrExpr : public Expr<TsrExpr<Array, Alias>> {
         TiledArray::expressions::is_aliased<D>::value,
         "no_alias() expressions are not allowed on the right-hand side of "
         "the assignment operator.");
-    return operator=(SubtExpr<TsrExpr_, D>(*this, other.derived()));
+    if (this->array().is_initialized()) {
+      return operator=(SubtExpr<TsrExpr_, D>(*this, other.derived()));
+    } else {
+      return operator=(-1 * other.derived());
+    }
   }
 
   /// Expression multiply-assignment operator
@@ -173,7 +181,11 @@ class TsrExpr : public Expr<TsrExpr<Array, Alias>> {
         TiledArray::expressions::is_aliased<D>::value,
         "no_alias() expressions are not allowed on the right-hand side of "
         "the assignment operator.");
-    return operator=(MultExpr<TsrExpr_, D>(*this, other.derived()));
+    if (this->array().is_initialized()) {
+      return operator=(MultExpr<TsrExpr_, D>(*this, other.derived()));
+    } else {
+      return operator=(0 * other.derived());
+    }
   }
 
   /// Array accessor

diff --git a/tests/expressions_impl.h b/tests/expressions_impl.h
@@ -1049,6 +1049,23 @@ BOOST_FIXTURE_TEST_CASE_TEMPLATE(add_to, F, Fixtures, F) {
       BOOST_CHECK(a.is_zero(i) && b.is_zero(i));
     }
   }
+
+  // try += with null this
+  std::decay_t<decltype(c)> d;
+  BOOST_REQUIRE_NO_THROW(d("a,b,c") += b("a,b,c"));
+
+  for (std::size_t i = 0ul; i < d.size(); ++i) {
+    if (!d.is_zero(i)) {
+      BOOST_CHECK(!b.is_zero(i));
+      auto b_tile = b.find(i).get();
+      auto d_tile = d.find(i).get();
+
+      for (std::size_t j = 0ul; j < d_tile.size(); ++j)
+        BOOST_CHECK_EQUAL(d_tile[j], b_tile[j]);
+    } else {
+      BOOST_CHECK(b.is_zero(i));
+    }
+  }
 }
 
 BOOST_FIXTURE_TEST_CASE_TEMPLATE(add_permute, F, Fixtures, F) {
@@ -1350,6 +1367,23 @@ BOOST_FIXTURE_TEST_CASE_TEMPLATE(subt_to, F, Fixtures, F) {
       BOOST_CHECK(a.is_zero(i) && b.is_zero(i));
     }
   }
+
+  // try -= with null this
+  std::decay_t<decltype(c)> d;
+  BOOST_REQUIRE_NO_THROW(d("a,b,c") -= b("a,b,c"));
+
+  for (std::size_t i = 0ul; i < d.size(); ++i) {
+    if (!d.is_zero(i)) {
+      BOOST_CHECK(!b.is_zero(i));
+      auto b_tile = b.find(i).get();
+      auto d_tile = d.find(i).get();
+
+      for (std::size_t j = 0ul; j < d_tile.size(); ++j)
+        BOOST_CHECK_EQUAL(d_tile[j], -b_tile[j]);
+    } else {
+      BOOST_CHECK(b.is_zero(i));
+    }
+  }
 }
 
 BOOST_FIXTURE_TEST_CASE_TEMPLATE(sub_permute, F, Fixtures, F) {
@@ -1679,6 +1713,14 @@ BOOST_FIXTURE_TEST_CASE_TEMPLATE(mult_to, F, Fixtures, F) {
       BOOST_CHECK(a.is_zero(i) || b.is_zero(i));
     }
   }
+
+  // try *= with null this
+  std::decay_t<decltype(c)> d;
+  BOOST_REQUIRE_NO_THROW(d("a,b,c") *= b("a,b,c"));
+
+  for (std::size_t i = 0ul; i < d.size(); ++i) {
+    BOOST_CHECK(d.is_zero(i));
+  }
 }
 
 BOOST_FIXTURE_TEST_CASE_TEMPLATE(scale_mult, F, Fixtures, F) {