[FFI] Replace Arg2Str with a more powerful for_each (#18117)

[FFI] Replace Arg2Str with a more powerful for_each

Author: PragmaTwice

ffi/include/tvm/ffi/base_details.h

@@ -29,6 +29,7 @@
 #include <tvm/ffi/endian.h>
 
 #include <cstddef>
+#include <type_traits>
 #include <utility>
 
 #if defined(_MSC_VER)
@@ -135,14 +136,32 @@ namespace tvm {
 namespace ffi {
 namespace details {
 
+// a dependent-name version of false, for static_assert
+template <typename>
+inline constexpr bool always_false = false;
+
 // for each iterator
 struct for_each_dispatcher {
   template <typename F, typename... Args, size_t... I>
   static void run(std::index_sequence<I...>, const F& f, Args&&... args) {  // NOLINT(*)
-    (f(I, std::forward<Args>(args)), ...);
+    if constexpr (std::conjunction_v<
+                      std::is_invocable<F, std::integral_constant<size_t, I>, Args>...>) {
+      (f(std::integral_constant<size_t, I>{}, std::forward<Args>(args)), ...);
+    } else if constexpr (std::conjunction_v<std::is_invocable<F, size_t, Args>...>) {
+      (f(I, std::forward<Args>(args)), ...);
+    } else if constexpr (std::conjunction_v<std::is_invocable<F, Args>...>) {
+      (f(std::forward<Args>(args)), ...);
+    } else {
+      static_assert(always_false<F>, "The function is not invocable with the provided arguments");
+    }
   }
 };
 
+// Three kinds of function F are acceptable in `for_each`:
+// 1. F(size_t, Arg): argument with its index
+// 2. F(Arg): just the argument
+// 3. F(std::integral_constant<size_t, I>, Arg): argument with its constexpr index
+// The third one can make the index available in template arguments and `if constexpr`.
 template <typename F, typename... Args>
 void for_each(const F& f, Args&&... args) {  // NOLINT(*)
   for_each_dispatcher::run(std::index_sequence_for<Args...>{}, f, std::forward<Args>(args)...);

ffi/include/tvm/ffi/function_details.h

@@ -36,23 +36,6 @@ namespace tvm {
 namespace ffi {
 namespace details {
 
-template <typename ArgType>
-struct Arg2Str {
-  template <size_t i>
-  TVM_FFI_INLINE static void Apply(std::ostream& os) {
-    using Arg = std::tuple_element_t<i, ArgType>;
-    if constexpr (i != 0) {
-      os << ", ";
-    }
-    os << i << ": " << Type2Str<Arg>::v();
-  }
-  template <size_t... I>
-  TVM_FFI_INLINE static void Run(std::ostream& os, std::index_sequence<I...>) {
-    using TExpander = int[];
-    (void)TExpander{0, (Apply<I>(os), 0)...};
-  }
-};
-
 template <typename T>
 static constexpr bool ArgSupported =
     (std::is_same_v<std::remove_const_t<std::remove_reference_t<T>>, Any> ||
@@ -78,10 +61,16 @@ struct FuncFunctorImpl {
 #endif
 
   TVM_FFI_INLINE static std::string Sig() {
-    using IdxSeq = std::make_index_sequence<sizeof...(Args)>;
     std::ostringstream ss;
     ss << "(";
-    Arg2Str<std::tuple<Args...>>::Run(ss, IdxSeq{});
+    for_each(
+        [&ss](auto i, const auto& v) {
+          if constexpr (i() != 0) {
+            ss << ", ";
+          }
+          ss << i() << ": " << v;
+        },
+        Type2Str<Args>::v()...);
     ss << ") -> " << Type2Str<R>::v();
     return ss.str();
   }