From 3427e31331677ca826c5588c87924214f7e5c54b Mon Sep 17 00:00:00 2001
From: Ville Voutilainen <ville.voutilainen@gmail.com>
Date: Fri, 9 Oct 2020 20:47:01 +0300
Subject: [PATCH] libstdc++: Diagnose visitors with different return types
 [PR95904]

libstdc++-v3/ChangeLog:

	PR libstdc++/95904
	* include/std/variant (__deduce_visit_result): Add a nested ::type.
	(__gen_vtable_impl</*base case*/>::_S_apply):
	Check the visitor return type.
	(__same_types): New.
	(__check_visitor_result): Likewise.
	(__check_visitor_results): Likewise.
	(visit(_Visitor&&, _Variants&&...)): Use __check_visitor_results
	in case we're visiting just one variant.
	* testsuite/20_util/variant/visit_neg.cc: Adjust.
---
 libstdc++-v3/include/std/variant              | 67 +++++++++++++++++--
 .../testsuite/20_util/variant/visit_neg.cc    |  2 +-
 2 files changed, 64 insertions(+), 5 deletions(-)

diff --git a/libstdc++-v3/include/std/variant b/libstdc++-v3/include/std/variant
index dd8847cf829..b32e564fd41 100644
--- a/libstdc++-v3/include/std/variant
+++ b/libstdc++-v3/include/std/variant
@@ -182,7 +182,7 @@ namespace __variant
   // used for raw visitation with indices passed in
   struct __variant_idx_cookie { using type = __variant_idx_cookie; };
   // Used to enable deduction (and same-type checking) for std::visit:
-  template<typename> struct __deduce_visit_result { };
+  template<typename _Tp> struct __deduce_visit_result { using type = _Tp; };
 
   // Visit variants that might be valueless.
   template<typename _Visitor, typename... _Variants>
@@ -1017,7 +1017,26 @@ namespace __variant
 
       static constexpr auto
       _S_apply()
-      { return _Array_type{&__visit_invoke}; }
+      {
+	if constexpr (_Array_type::__result_is_deduced::value)
+	  {
+	    constexpr bool __visit_ret_type_mismatch =
+	      !is_same_v<typename _Result_type::type,
+			 decltype(__visit_invoke(std::declval<_Visitor>(),
+				    std::declval<_Variants>()...))>;
+	    if constexpr (__visit_ret_type_mismatch)
+	      {
+		static_assert(!__visit_ret_type_mismatch,
+		  "std::visit requires the visitor to have the same "
+		  "return type for all alternatives of a variant");
+		return __nonesuch{};
+	      }
+	    else
+	      return _Array_type{&__visit_invoke};
+	  }
+	else
+	  return _Array_type{&__visit_invoke};
+      }
     };
 
   template<typename _Result_type, typename _Visitor, typename... _Variants>
@@ -1692,6 +1711,26 @@ namespace __variant
 			   std::forward<_Variants>(__variants)...);
     }
 
+  template<typename _Tp, typename... _Types>
+     constexpr inline bool __same_types = (is_same_v<_Tp, _Types> && ...);
+
+  template <unsigned long int _Idx, typename _Visitor, typename _Variant>
+    decltype(auto)
+    __check_visitor_result(_Visitor&& __vis, _Variant&& __variant)
+    {
+      return std::__invoke(std::forward<_Visitor>(__vis),
+			   std::get<_Idx>(std::forward<_Variant>(__variant)));
+    }
+
+  template <typename _Visitor, typename _Variant, unsigned long int... _Idxs>
+    constexpr bool __check_visitor_results(std::index_sequence<_Idxs...>)
+    {
+      return __same_types<decltype(__check_visitor_result<_Idxs>(
+	std::declval<_Visitor>(),
+	std::declval<_Variant>()))...>;
+    }
+
+
   template<typename _Visitor, typename... _Variants>
     constexpr decltype(auto)
     visit(_Visitor&& __visitor, _Variants&&... __variants)
@@ -1704,8 +1743,28 @@ namespace __variant
 
       using _Tag = __detail::__variant::__deduce_visit_result<_Result_type>;
 
-      return std::__do_visit<_Tag>(std::forward<_Visitor>(__visitor),
-				   std::forward<_Variants>(__variants)...);
+      if constexpr (sizeof...(_Variants) == 1)
+	{
+	  constexpr bool __visit_rettypes_match =
+	    __check_visitor_results<_Visitor, _Variants...>(
+	      std::make_index_sequence<
+	        std::variant_size<remove_reference_t<_Variants>...>::value>());
+	  if constexpr (!__visit_rettypes_match)
+	    {
+	      static_assert(__visit_rettypes_match,
+			  "std::visit requires the visitor to have the same "
+			  "return type for all alternatives of a variant");
+	      return;
+	    }
+	  else
+	    return std::__do_visit<_Tag>(
+	      std::forward<_Visitor>(__visitor),
+	      std::forward<_Variants>(__variants)...);
+	}
+      else
+	return std::__do_visit<_Tag>(
+	  std::forward<_Visitor>(__visitor),
+	  std::forward<_Variants>(__variants)...);
     }
 
 #if __cplusplus > 201703L
diff --git a/libstdc++-v3/testsuite/20_util/variant/visit_neg.cc b/libstdc++-v3/testsuite/20_util/variant/visit_neg.cc
index 6279dec5aa2..748eb21c1ad 100644
--- a/libstdc++-v3/testsuite/20_util/variant/visit_neg.cc
+++ b/libstdc++-v3/testsuite/20_util/variant/visit_neg.cc
@@ -21,7 +21,7 @@
 #include <variant>
 #include <testsuite_hooks.h>
 
-// { dg-error "invalid conversion" "" { target *-*-* } 0 }
+// { dg-error "same return type for all alternatives" "" { target *-*-* } 0 }
 // { dg-prune-output "in 'constexpr' expansion" }
 
 void
-- 
2.30.2