WIP Add optional TaskStatus for StateDescriptor callbacks

src/CMakeLists.txt

@@ -277,6 +277,7 @@ add_library(parthenon
   utils/communication_buffer.hpp
   utils/cleantypes.hpp
   utils/concepts_lite.hpp
+  utils/default_return_function.hpp
   utils/error_checking.cpp
   utils/error_checking.hpp
   utils/hash.hpp

src/interface/state_descriptor.hpp

@@ -33,10 +33,10 @@
 #include "outputs/output_parameters.hpp"
 #include "parameter_input.hpp"
 #include "prolong_restrict/prolong_restrict.hpp"
+#include "utils/default_return_function.hpp"
 #include "utils/error_checking.hpp"
 
 namespace parthenon {
-
 // Forward declarations
 template <typename T>
 class MeshBlockData;
@@ -350,29 +350,37 @@ class StateDescriptor {
 
   bool FlagsPresent(std::vector<MetadataFlag> const &flags, bool matchAny = false);
 
-  void PreCommFillDerived(MeshBlockData<Real> *rc) const {
-    if (PreCommFillDerivedBlock != nullptr) PreCommFillDerivedBlock(rc);
+  TaskStatus PreCommFillDerived(MeshBlockData<Real> *rc) const {
+    if (PreCommFillDerivedBlock != nullptr) return PreCommFillDerivedBlock(rc);
+    return TaskStatus::complete;
   }
-  void PreCommFillDerived(MeshData<Real> *rc) const {
-    if (PreCommFillDerivedMesh != nullptr) PreCommFillDerivedMesh(rc);
+  TaskStatus PreCommFillDerived(MeshData<Real> *rc) const {
+    if (PreCommFillDerivedMesh != nullptr) return PreCommFillDerivedMesh(rc);
+    return TaskStatus::complete;
   }
-  void PreFillDerived(MeshBlockData<Real> *rc) const {
-    if (PreFillDerivedBlock != nullptr) PreFillDerivedBlock(rc);
+  TaskStatus PreFillDerived(MeshBlockData<Real> *rc) const {
+    if (PreFillDerivedBlock != nullptr) return PreFillDerivedBlock(rc);
+    return TaskStatus::complete;
   }
-  void PreFillDerived(MeshData<Real> *rc) const {
-    if (PreFillDerivedMesh != nullptr) PreFillDerivedMesh(rc);
+  TaskStatus PreFillDerived(MeshData<Real> *rc) const {
+    if (PreFillDerivedMesh != nullptr) return PreFillDerivedMesh(rc);
+    return TaskStatus::complete;
   }
-  void PostFillDerived(MeshBlockData<Real> *rc) const {
-    if (PostFillDerivedBlock != nullptr) PostFillDerivedBlock(rc);
+  TaskStatus PostFillDerived(MeshBlockData<Real> *rc) const {
+    if (PostFillDerivedBlock != nullptr) return PostFillDerivedBlock(rc);
+    return TaskStatus::complete;
   }
-  void PostFillDerived(MeshData<Real> *rc) const {
-    if (PostFillDerivedMesh != nullptr) PostFillDerivedMesh(rc);
+  TaskStatus PostFillDerived(MeshData<Real> *rc) const {
+    if (PostFillDerivedMesh != nullptr) return PostFillDerivedMesh(rc);
+    return TaskStatus::complete;
   }
-  void FillDerived(MeshBlockData<Real> *rc) const {
-    if (FillDerivedBlock != nullptr) FillDerivedBlock(rc);
+  TaskStatus FillDerived(MeshBlockData<Real> *rc) const {
+    if (FillDerivedBlock != nullptr) return FillDerivedBlock(rc);
+    return TaskStatus::complete;
   }
-  void FillDerived(MeshData<Real> *rc) const {
-    if (FillDerivedMesh != nullptr) FillDerivedMesh(rc);
+  TaskStatus FillDerived(MeshData<Real> *rc) const {
+    if (FillDerivedMesh != nullptr) return FillDerivedMesh(rc);
+    return TaskStatus::complete;
   }
 
   void PreStepDiagnostics(SimTime const &simtime, MeshData<Real> *rc) const {
@@ -425,14 +433,23 @@ class StateDescriptor {
 
   std::vector<std::shared_ptr<AMRCriteria>> amr_criteria;
 
-  std::function<void(MeshBlockData<Real> *rc)> PreCommFillDerivedBlock = nullptr;
-  std::function<void(MeshData<Real> *rc)> PreCommFillDerivedMesh = nullptr;
-  std::function<void(MeshBlockData<Real> *rc)> PreFillDerivedBlock = nullptr;
-  std::function<void(MeshData<Real> *rc)> PreFillDerivedMesh = nullptr;
-  std::function<void(MeshBlockData<Real> *rc)> PostFillDerivedBlock = nullptr;
-  std::function<void(MeshData<Real> *rc)> PostFillDerivedMesh = nullptr;
-  std::function<void(MeshBlockData<Real> *rc)> FillDerivedBlock = nullptr;
-  std::function<void(MeshData<Real> *rc)> FillDerivedMesh = nullptr;
+  // function_t behaves like a std::function<return_t(MeshData<Real>*)>. If the
+  // function it points to returns a TaskStatus, this is passed back to the caller.
+  // If return_t is any other type, calling function_t will return TaskStatus::complete.
+  using function_t =
+      DefaultReturnFunction<TaskStatus, TaskStatus::complete, MeshData<Real> *>;
+  using function_block_t =
+      DefaultReturnFunction<TaskStatus, TaskStatus::complete, MeshBlockData<Real> *>;
+
+  function_block_t PreCommFillDerivedBlock;
+  function_t PreCommFillDerivedMesh;
+  function_block_t PreFillDerivedBlock;
+  function_t PreFillDerivedMesh;
+  function_block_t PostFillDerivedBlock;
+  function_t PostFillDerivedMesh;
+  function_block_t FillDerivedBlock;
+  function_t FillDerivedMesh;
+
   std::function<void(Mesh *, ParameterInput *, SimTime &)> UserWorkBeforeLoopMesh =
       nullptr;
   std::function<void(Mesh *, ParameterInput *, SimTime &)> UserWorkBeforeOutputMesh =

src/interface/update.hpp

@@ -282,7 +282,8 @@ TaskStatus PreCommFillDerived(T *rc) {
   PARTHENON_INSTRUMENT
   auto pm = rc->GetParentPointer();
   for (const auto &pkg : pm->packages.AllPackages()) {
-    pkg.second->PreCommFillDerived(rc);
+    auto status = pkg.second->PreCommFillDerived(rc);
+    if (status != TaskStatus::complete) return status;
   }
   return TaskStatus::complete;
 }
@@ -294,19 +295,22 @@ TaskStatus FillDerived(T *rc) {
   { // PreFillDerived region
     PARTHENON_INSTRUMENT
     for (const auto &pkg : pm->packages.AllPackages()) {
-      pkg.second->PreFillDerived(rc);
+      auto status = pkg.second->PreFillDerived(rc);
+      if (status != TaskStatus::complete) return status;
     }
   } // PreFillDerived region
   { // FillDerived region
     PARTHENON_INSTRUMENT
     for (const auto &pkg : pm->packages.AllPackages()) {
-      pkg.second->FillDerived(rc);
+      auto status = pkg.second->FillDerived(rc);
+      if (status != TaskStatus::complete) return status;
     }
   } // FillDerived region
   { // PostFillDerived region
     PARTHENON_INSTRUMENT
     for (const auto &pkg : pm->packages.AllPackages()) {
-      pkg.second->PostFillDerived(rc);
+      auto status = pkg.second->PostFillDerived(rc);
+      if (status != TaskStatus::complete) return status;
     }
   } // PostFillDerived region
   return TaskStatus::complete;

src/utils/default_return_function.hpp

@@ -0,0 +1,62 @@
+//========================================================================================
+// (C) (or copyright) 2023. Triad National Security, LLC. All rights reserved.
+//
+// This program was produced under U.S. Government contract 89233218CNA000001 for Los
+// Alamos National Laboratory (LANL), which is operated by Triad National Security, LLC
+// for the U.S. Department of Energy/National Nuclear Security Administration. All rights
+// in the program are reserved by Triad National Security, LLC, and the U.S. Department
+// of Energy/National Nuclear Security Administration. The Government is granted for
+// itself and others acting on its behalf a nonexclusive, paid-up, irrevocable worldwide
+// license in this material to reproduce, prepare derivative works, distribute copies to
+// the public, perform publicly and display publicly, and to permit others to do so.
+//========================================================================================
+
+#ifndef UTILS_DEFAULT_RETURN_FUNCTION_HPP_
+#define UTILS_DEFAULT_RETURN_FUNCTION_HPP_
+
+#include <type_traits>
+#include <utility>
+
+namespace parthenon {
+template <class return_t, return_t default_ret, class... args_t>
+class DefaultReturnFunction {
+ public:
+  DefaultReturnFunction() : func_(nullptr) {}
+  explicit DefaultReturnFunction(std::nullptr_t) : func_(nullptr) {}
+
+  template <class F, REQUIRES(std::is_invocable_v<std::decay_t<F>, args_t...>)>
+  explicit DefaultReturnFunction(F &&f) {
+    assign(std::forward<F>(f));
+  }
+
+  template <class F, REQUIRES(std::is_invocable_v<std::decay_t<F>, args_t...>)>
+  DefaultReturnFunction &operator=(F &&f) {
+    assign(std::forward<F>(f));
+    return *this;
+  }
+
+  TaskStatus operator()(args_t... args) const { return func_(args...); }
+  bool operator==(std::nullptr_t) const { return !func_; }
+  bool operator!=(std::nullptr_t) const { return static_cast<bool>(func_); }
+  explicit operator bool() const { return static_cast<bool>(func_); }
+
+ private:
+  std::function<return_t(args_t...)> func_;
+
+  template <class F>
+  void assign(F &&f) {
+    if constexpr (std::is_same_v<std::invoke_result_t<std::decay_t<F>, args_t...>,
+                                 return_t>) {
+      func_ = f;
+    } else {
+      func_ = [f](args_t... args) {
+        f(args...);
+        return default_ret;
+      };
+    }
+  }
+};
+
+} // namespace parthenon
+
+#endif // UTILS_DEFAULT_RETURN_FUNCTION_HPP_