From f514f2f49ba6b206a46b2ec19997e79111cb1277 Mon Sep 17 00:00:00 2001
From: Marcel Koch <marcel.koch@kit.edu>
Date: Wed, 3 Apr 2024 10:48:10 +0200
Subject: [PATCH] [core] deduce constness in run

---
 core/base/block_operator.cpp  |  5 ++--
 core/base/dispatch_helper.hpp | 49 ++++++++++++++++++++++++++++++-----
 core/matrix/dense.cpp         |  4 +--
 core/matrix/permutation.cpp   |  4 +--
 core/matrix/row_gatherer.cpp  |  6 ++---
 5 files changed, 50 insertions(+), 18 deletions(-)

diff --git a/core/base/block_operator.cpp b/core/base/block_operator.cpp
index 1efc5e3b152..b8190bad02d 100644
--- a/core/base/block_operator.cpp
+++ b/core/base/block_operator.cpp
@@ -22,9 +22,8 @@ namespace {
 template <typename Fn>
 auto dispatch_dense(Fn&& fn, LinOp* v)
 {
-    return run<matrix::Dense<float>*, matrix::Dense<double>*,
-               matrix::Dense<std::complex<float>>*,
-               matrix::Dense<std::complex<double>>*>(v, std::forward<Fn>(fn));
+    return run<matrix::Dense, float, double, std::complex<float>,
+               std::complex<double>>(v, std::forward<Fn>(fn));
 }
 
 
diff --git a/core/base/dispatch_helper.hpp b/core/base/dispatch_helper.hpp
index 35422a75fd2..8fb104c084e 100644
--- a/core/base/dispatch_helper.hpp
+++ b/core/base/dispatch_helper.hpp
@@ -16,6 +16,12 @@ namespace gko {
 namespace detail {
 
 
+template <typename T, typename MaybeConstU>
+using with_same_constness_t = std::conditional_t<
+    std::is_const<typename std::remove_reference_t<MaybeConstU>>::value,
+    const T, T>;
+
+
 /**
  *
  * @copydoc run<typename ReturnType, typename K, typename... Types, typename T,
@@ -24,7 +30,7 @@ namespace detail {
  * @note this is the end case
  */
 template <typename ReturnType, typename T, typename Func, typename... Args>
-ReturnType run_impl(T obj, Func&&, Args&&...)
+ReturnType run_impl(T* obj, Func&&, Args&&...)
 {
     GKO_NOT_SUPPORTED(obj);
 }
@@ -37,9 +43,9 @@ ReturnType run_impl(T obj, Func&&, Args&&...)
  */
 template <typename ReturnType, typename K, typename... Types, typename T,
           typename Func, typename... Args>
-ReturnType run_impl(T obj, Func&& f, Args&&... args)
+ReturnType run_impl(T* obj, Func&& f, Args&&... args)
 {
-    if (auto dobj = dynamic_cast<K>(obj)) {
+    if (auto dobj = dynamic_cast<with_same_constness_t<K, T>*>(obj)) {
         return f(dobj, std::forward<Args>(args)...);
     } else {
         return run_impl<ReturnType, Types...>(obj, std::forward<Func>(f),
@@ -140,13 +146,16 @@ ReturnType run_impl(T obj, Func&& f, Args&&... args)
  */
 template <typename K, typename... Types, typename T, typename Func,
           typename... Args>
-auto run(T obj, Func&& f, Args&&... args)
+auto run(T* obj, Func&& f, Args&&... args)
 {
 #if __cplusplus < 201703L
     // result_of_t is deprecated in C++17
-    using ReturnType = std::result_of_t<Func(K, Args...)>;
+    using ReturnType =
+        std::result_of_t<Func(detail::with_same_constness_t<K, T>*, Args...)>;
 #else
-    using ReturnType = std::invoke_result_t<Func, K, Args...>;
+    using ReturnType =
+        std::invoke_result_t<Func, detail::with_same_constness_t<K, T>*,
+                             Args...>;
 #endif
     return detail::run_impl<ReturnType, K, Types...>(
         obj, std::forward<Func>(f), std::forward<Args>(args)...);
@@ -188,6 +197,34 @@ auto run(std::shared_ptr<T> obj, Func&& f, Args&&... args)
 }
 
 
+/**
+ * run uses template to go through the list and select the valid
+ * template and run it.
+ *
+ * @tparam Base  the Base class with one template
+ * @tparam K  the current template type of B. pointer of const Base<K> is tried
+ *            in the conversion.
+ * @tparam ...Types  the other types will be tried in the conversion if K fails
+ * @tparam T  the type of input object waiting converted
+ * @tparam Func  the function will run if the object can be converted to pointer
+ *               of const Base<K>
+ * @tparam ...Args  the additional arguments for the Func
+ *
+ * @param obj  the input object waiting converted
+ * @param f  the function will run if obj can be converted successfully
+ * @param args  the additional arguments for the function
+ *
+ * @return  the result of f invoked with obj cast to the first matching type
+ */
+template <template <class> class K, typename... Types, typename T,
+          typename Func, typename... Args>
+auto run(T* obj, Func&& f, Args&&... args)
+{
+    return run<K<Types>...>(obj, std::forward<Func>(f),
+                            std::forward<Args>(args)...);
+}
+
+
 /**
  * run uses template to go through the list and select the valid
  * template and run it.
diff --git a/core/matrix/dense.cpp b/core/matrix/dense.cpp
index 39f3f2fcb51..d1f3da8e166 100644
--- a/core/matrix/dense.cpp
+++ b/core/matrix/dense.cpp
@@ -1522,9 +1522,7 @@ template <typename ValueType, typename Function>
 void gather_mixed_real_complex(Function fn, LinOp* out)
 {
 #ifdef GINKGO_MIXED_PRECISION
-    using fst_type = matrix::Dense<ValueType>;
-    using snd_type = matrix::Dense<next_precision<ValueType>>;
-    run<fst_type*, snd_type*>(out, fn);
+    run<matrix::Dense, ValueType, next_precision<ValueType>>(out, fn);
 #else
     precision_dispatch<ValueType>(fn, out);
 #endif
diff --git a/core/matrix/permutation.cpp b/core/matrix/permutation.cpp
index b596fdae98c..84989c7eddb 100644
--- a/core/matrix/permutation.cpp
+++ b/core/matrix/permutation.cpp
@@ -269,8 +269,8 @@ void dispatch_dense(const LinOp* op, Functor fn)
 {
     using matrix::Dense;
     using std::complex;
-    run<const Dense<double>*, const Dense<float>*,
-        const Dense<complex<double>>*, const Dense<complex<float>>*>(op, fn);
+    run<Dense, double, float, std::complex<double>, std::complex<float>>(op,
+                                                                         fn);
 }
 
 
diff --git a/core/matrix/row_gatherer.cpp b/core/matrix/row_gatherer.cpp
index 09a067cf35e..0f570cda410 100644
--- a/core/matrix/row_gatherer.cpp
+++ b/core/matrix/row_gatherer.cpp
@@ -66,8 +66,7 @@ RowGatherer<IndexType>::create_const(
 template <typename IndexType>
 void RowGatherer<IndexType>::apply_impl(const LinOp* in, LinOp* out) const
 {
-    run<const Dense<float>*, const Dense<double>*,
-        const Dense<std::complex<float>>*, const Dense<std::complex<double>>*>(
+    run<Dense, float, double, std::complex<float>, std::complex<double>>(
         in, [&](auto gather) { gather->row_gather(&row_idxs_, out); });
 }
 
@@ -75,8 +74,7 @@ template <typename IndexType>
 void RowGatherer<IndexType>::apply_impl(const LinOp* alpha, const LinOp* in,
                                         const LinOp* beta, LinOp* out) const
 {
-    run<const Dense<float>*, const Dense<double>*,
-        const Dense<std::complex<float>>*, const Dense<std::complex<double>>*>(
+    run<Dense, float, double, std::complex<float>, std::complex<double>>(
         in,
         [&](auto gather) { gather->row_gather(alpha, &row_idxs_, beta, out); });
 }